package GraphAlgorithm;

import java.io.BufferedWriter;

import java.io.File;

import java.io.FileWriter;

import java.io.IOException;

import java.util.ArrayList;

import java.util.Collections;

import java.util.Comparator;

import java.util.HashMap;

import java.util.LinkedHashMap;

import java.util.LinkedList;

import java.util.List;

import java.util.Map;

import java.util.Set;

import network.CommGeneticsLoader;

import network.DisGraph;

import network.SparseMatrix;

import network.SparseVector;

import util.DCNOMIMUMLSIDmap;

import util.FileToList;

/**

 * Implement random walk with restart algorithm using sparse matrix data structure

 * @author zhengc

 *

 */

public class RandomWalk {

//  static final double P = 0.15;

//  static final double P = 0.15;

    static final double err = 0.000001;

    private Map<String, Integer> generank;

    /**

     * Constructs a new RandomWalk object from a graph and seed vector

     * @param cgg a CommGeneticsGraph 

     * @param E the seed vector

     * @param p restart rate

     */

    public RandomWalk(DisGraph cgg, SparseVector E, double p) {

        Map<String, Double> genescore = new HashMap<>();

        int size = cgg.getNodes();

        SparseVector x = new SparseVector(size);

        SparseMatrix A = cgg.getNet();

        A = normalizeMatrix(A);

        int iter = 0;

        SparseVector last_x = E;

        while((last_x.plus(x.scale(-1)).norm()) > err) {

            last_x = x;

            x = A.times(x).scale(1-p);

            x = x.plus(E.scale(p));

//          System.out.print(x);

            iter++;

//          System.out.println("Finish iteration " + iter + ":");

        }

        System.out.println("Total iteration " + iter);

        genescore = svToMap(x, E);

        generank = scoreToRank(genescore);

    }

    /**

     * Get the gene-rank pairs from random walk with restart

     * @return the gene-rank pairs from random walk with restart

     */

    public  Map<String, Integer> getRWRank() {

        return generank;

    }

    /**

     * Generate column/left stochastic matrix by normalization

     * @param comMatrix

     * @return normalized SparseMatrix

     */

    private static SparseMatrix normalizeMatrix(SparseMatrix comMatrix) {

        int n = comMatrix.size();

        double[] total_w = new double[n]; //Total weight of a column

        SparseMatrix A = new SparseMatrix(n);

        for (int j = 0; j < n; j++){

            total_w[j] = 0;

            for (int i = 0; i < n; i++){

                total_w[j] += comMatrix.get(i, j);

            }

        }

        //normalize the adjacency matrix

        for (int i = 0; i < n; i++){

            for (int j = 0; j < n; j++){

                double norm_w = comMatrix.get(i, j);

                norm_w = norm_w * 1.0 / total_w[j];

                A.put(i, j, norm_w);

            }

        }

        return A;

    }

    /**

     * Create a seed from a list

     * @param nodesids a list of seed nodes

     * @return the seed vector

     * @throws IOException

     */

    public static SparseVector createSeedVector(List<String> nodesids) throws IOException{

        int n = CommGeneticsLoader.entry_index.size();

        SparseVector E = new SparseVector(n);

        for (String id : nodesids) {

            if (CommGeneticsLoader.entry_index.containsKey(id)) {

                int index = CommGeneticsLoader.entry_index.get(id);

//              E.put(index, 1);

                E.put(index, 1.0 / nodesids.size());

            } else {

                System.out.println(id + " is NOT included in our network!");

            }

        }

        return E;

    }

    public static SparseVector createSeedVectorFromIdx(List<Integer> nodeidxs) {

        int n = CommGeneticsLoader.entry_index.size();

        SparseVector E = new SparseVector(n);

        for (int idx : nodeidxs) {

            if (idx > n - 1) {

                System.out.println(idx + " is out of network!");

            } else {

                E.put(idx, 1.0 / nodeidxs.size());;

            }

        }

        return E;

    }

    /**

     * Convert a SparseVector from random walk into a HashMap reversely ordered by value

     * @param sv the SparseVector representing gene rank score result

     * @param seeds_vector the seed vector

     * @return a HashMap reversely ordered by value 

     */

    private static Map<String, Double> svToMap(SparseVector sv, SparseVector seeds_vector) {

        Map<String, Double> genemap = new HashMap<String,Double>();

//      System.out.println(CommGeneticsLoader.num_disease);

//      System.out.println(CommGeneticsLoader.entry_list.size());

        for (int i = CommGeneticsLoader.num_disease; i < CommGeneticsLoader.entry_list.size(); i++){

//          double val_seed = seeds_vector.get(i);

            Set<Integer> seeds_key = seeds_vector.getKeys();

            if (!seeds_key.contains(i)) { // Test if the node is one of seeds

                double val = sv.get(i);

//              System.out.println(index_entry.get(num_disease));

                if (val != 0){

                    String gene = CommGeneticsLoader.index_entry.get(i);

                    genemap.put(gene, val);

                }

            }

        }

        genemap = util.CollectionsEx.sortByValueR(genemap);

        return genemap;

    }

    /**

     * Convert a gene rank score into a gene rank HashMap reversely ordered by value

     * @param sortedgenemap a reversely sorted gene score map

     * @return a gene rank  HashMap

     */

    private static Map<String, Integer> scoreToRank(Map<String, Double> sortedgenemap) {

        Map<String, Integer> generank = new HashMap<String,Integer>();

        int rank = 0;

        for (String s : sortedgenemap.keySet()) {

//          System.out.println(sortedgenemap.get(s));

            rank++;

            generank.put(s, rank);

            generank = util.CollectionsEx.sortByValue(generank);

        }

        return generank;

    }

//  System.out.println(genemap.size());

    /**

     * Write gene rank scores into CSV file, file format is "Rank, Gene, Score"

     * @param sortedgenemap HashMap sorted by descend values

     * @param rankfile file to be written

     * @throws IOException

     */

    public static void saveScorefile(Map<String, Double> sortedgenemap, String rankfile) throws IOException {

        BufferedWriter bw = new BufferedWriter(new FileWriter(new File(rankfile)));

        int rank = 0;

        bw.write("Rank" + "," + "Gene" + "," + "Score" + "," + "\n");

        for (String s : sortedgenemap.keySet()) {

//          System.out.println(sortedgenemap.get(s));

            Double score = sortedgenemap.get(s);

            rank++;

            bw.write(rank + "," + s + "," + score + "," + "\n");

        }

        bw.close();

    }

    /**

     * Write gene rank file into a CSV file, file format is "Rank, Gene"

     * @param generank gene rank HashMap

     * @param rankfile the file to be written

     * @throws IOException

     */

    public static void saveRankfile(Map<String, Integer> generank, String rankfile) throws IOException {

        BufferedWriter bw = new BufferedWriter(new FileWriter(new File(rankfile)));

        bw.write("Rank" + "," + "Gene" + ","  + "\n");

        for (String s : generank.keySet()) {

//          System.out.println(sortedgenemap.get(s));

            int rank = generank.get(s);

            bw.write(rank + "," + s + "," + "\n");

        }

        bw.close();

    }

    public static void main(String[] args) throws IOException {

        /*

         * Create additional DCN_OMIM map for give disease

         */

        String dcnmapfile = "./data/term_umls_id_diso";

        String omimmapfile = "./data/OMIM_umls_id_diso";

        util.DCNOMIMUMLSIDmap.createDCNIdNameMap(dcnmapfile);

        util.DCNOMIMUMLSIDmap.createOMIMIdNameMap(omimmapfile);

        String DCN_dis = "dementia";

        String OMIM_pat = "alzheimer";

        DCNOMIMUMLSIDmap domim = new DCNOMIMUMLSIDmap(DCN_dis, OMIM_pat);

        Map<String, List<String>> dcn_omim = domim.getDCNOMIMUMLSIDmap();

        /*

         * Build the bipartite network

         */

        String commnetfile = "./results/fares_comm_net_lift_final_abbr.txt";

        String ppifile = "./data/gene_gene_string_cut.txt";

        String disgenefile = "./data/OMIM_disease_gene_umls_id_diso";

        DisGraph cgg = null;

        try {

            cgg = CommGeneticsLoader.createCommGeneticsGraph(commnetfile, ppifile, disgenefile, dcn_omim);

        } catch (IOException e) {

            // TODO Auto-generated catch block

            e.printStackTrace();

        }

        /*

         * Create seed vector

         */

        int n = cgg.getNodes();

        List<String> seeds = new ArrayList<>();

        String seed_dis1 = "dementia";

        String seed_dis1_umls = null;

        if (util.DCNOMIMUMLSIDmap.dcnnameidmap.containsKey(seed_dis1)) {

            seed_dis1_umls = util.DCNOMIMUMLSIDmap.dcnnameidmap.get(seed_dis1);

            System.out.println(seed_dis1 + ": " + seed_dis1_umls);

            seeds.add(seed_dis1_umls);

        } else {

            System.out.println(seed_dis1 + " is NOT included in our network!");

        }

        String seed_genes_file = "./data/AD_omim_genes.txt";

        List<String> seed_genes = FileToList.createListFromfile(seed_genes_file);

        seeds.addAll(seed_genes);

//      diseases.add(disease2);

        SparseVector seeds_vec = new SparseVector(n);

        try {

            seeds_vec = createSeedVector(seeds);

        } catch (IOException e) {

            // TODO Auto-generated catch block

            e.printStackTrace();

        }

        System.out.println("Seed vector: " + seeds_vec);

        /*

         * Random walk

         */

        RandomWalk rwrank = new RandomWalk(cgg, seeds_vec, 0.5);

        Map<String, Integer> result = rwrank.getRWRank();

        /*

         * Process result and write to file

         */

        String rankfile = "./results/AD_novel_genes.csv";

        try {

            saveRankfile(result, rankfile);

        } catch (IOException e) {

            // TODO Auto-generated catch block

            e.printStackTrace();

        }

    }

}