package org.genemania.dw.tools;

import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.TreeMap;

import org.genemania.dw.entity.ExtResource;
import org.genemania.dw.entity.Uniprot;
import org.genemania.dw.util.GenUtil;

/**
 * Container class plus matching service, currently serving Uniprot vs.
 * Entrez/xrefID mappings.
 * Allows simple caching for mapped identifiers and flushing them.
 * The indexUniprotByXref() method can be called less during init () with more
 * refactoring.
 * The ExtResource.LIST_ID is used to reference EntrezID. Alternatively, we
 * can use ExtResource.RES_ENTREZ, or some alternative CV.
 *
 * @author rashadbadrawi
 */

public class IdentifierMapperService {

    public static final String [] REFSEQ_RNA_PREFIX_ARR = {"NM", "NR", "XM",
    "XR"};
    public static final String [] REFSEQ_PRO_PREFIX_ARR = {"AP", "NP", "XP",
    "YP", "ZP"};
    public static final String [] REFSEQ_GENO_PREFIX_ARR = {"AC", "NC", "NG",
    "NT", "NW", "NZ", "NS"};
    public static final String LO_PROTEIN = "LO_PROTEIN";
    public static final String LO_GI_UNI = "LO_GI_UNI";
    
    private static PrintWriter log;
    private static boolean needSuppMappings = true;
    private static TreeMap <String, TreeMap <String, ExtResource>> uniMap =
                           new TreeMap <String, TreeMap <String, ExtResource>> ();
    //handy maps
    private static TreeMap <String, TreeMap <String, String>> GIUniMap =
                           new TreeMap <String, TreeMap <String, String>> ();
    private static TreeMap <String, TreeMap <String, String>> entrezUniMap =
                           new TreeMap <String, TreeMap <String, String>> ();
    //for GI leftovers, specific for BIND, and across all species.
    private static TreeMap <String, String> GIRefSeqProMap = new TreeMap <String, String> ();
    private static TreeMap <String, String> GIProMap = new TreeMap <String, String> ();
    private static TreeMap <String, String> GIRefSeqMrnaMap = new TreeMap <String, String> ();
    private static TreeMap <String, ArrayList <String>> GIUniSupMap =
                                new TreeMap <String, ArrayList <String>> ();
    //not flushed
    private static ArrayList <String> unmatchedGIUniList = new ArrayList <String> ();
    private static ArrayList <String> matchedGIUniList = new ArrayList <String> ();
    private static ArrayList <String> unmatchedEntrezUniList = new ArrayList <String> ();
    private static ArrayList <String> matchedEntrezUniList = new ArrayList <String> ();

    private IdentifierMapperService () {

        log = GenUtil.getDefaultLog ();
    }

    public static void init (String [] resArr, int [] speciesIDArr) throws Exception {

        GenUtil.validateNotNull(resArr);
        for (int i = 0; i < resArr.length; i++) {
            //System.out.println ("Init resource: " + resArr [i]);
            if (ExtResource.RES_UNIPROT.equals (resArr [i])) {
                TreeMap <String, ExtResource> uniSpeciesMap;
                boolean indexFlag = false;
                if (speciesIDArr != null) {
                    for (int j = 0; j < speciesIDArr.length; j++) {
                        if (uniMap.get (String.valueOf(speciesIDArr [j])) == null) {
                            System.out.println ("Processing: " + resArr [i] +
                                                " " + speciesIDArr [j]);
                            uniSpeciesMap = Uniprot.loadAllExt(speciesIDArr [j]);
                            //System.out.println ("Done: " + resArr [i] +
                            //                    " " + speciesIDArr [j]);
                            uniMap.put(String.valueOf (speciesIDArr [j]), uniSpeciesMap);
                            indexFlag = true;
                        }
                    }
                } else {
                    System.err.println ("Loading for Uniprot is by species only.");
                }
                if (indexFlag) {
                    System.out.println ("Indexing Uniprot by GI...");
                    indexUniprotByXref (speciesIDArr, ExtResource.LIST_GI);
                    System.out.println ("Indexing Uniprot by Entrez ID...");
                    indexUniprotByXref(speciesIDArr, ExtResource.LIST_ID);
                }
            } else {
                throw new IllegalArgumentException ("Unsupported resource type: " +
                                                    resArr [i]);
            }
            System.out.println ("Done init resource: " + resArr [i]);
        }
        if (IdentifierMapperService.needSuppMappings) {
            IdentifierMapperService.needSuppMappings = false;
            loadSupMappings();
        }
    }

    public static void flush (String [] resArr, int [] speciesIDArr) {

        GenUtil.validateNotNull(resArr);
        for (int i = 0; i < resArr.length; i++) {
            //System.out.println ("Flushing resource: " + resArr [i]);
            if (ExtResource.RES_UNIPROT.equals (resArr [i])) {
                if (speciesIDArr != null) {
                    for (int j = 0; j < speciesIDArr.length; j++) {
                        uniMap.remove(String.valueOf (speciesIDArr [j]));
                        GIUniMap.remove (String.valueOf (speciesIDArr [j]));
                        entrezUniMap.remove (String.valueOf (speciesIDArr [j]));
                    }
                } else {
                    uniMap.remove(GenUtil.ALL);
                    GIUniMap.remove (GenUtil.ALL);
                    entrezUniMap.remove(GenUtil.ALL);
                }
            } else {
                throw new IllegalArgumentException ("Unsupported resource type: " +
                                                    resArr [i]);
            }
            //System.out.println ("Done flushing resource: " + resArr [i]);
        }
    }

    private static void indexUniprotByXref (int [] speciesIDArr, String xrefType) {

        for (int i = 0; i < speciesIDArr.length; i++) {
            TreeMap <String, String> xrefUniSpeciesMap = new TreeMap <String, String> (); //local
            TreeMap <String, ExtResource> uniSpeciesMap = uniMap.get (String.valueOf (speciesIDArr [i]));
            Iterator iterator = uniSpeciesMap.keySet().iterator();
            while (iterator.hasNext()) {
                String uniID = (String)iterator.next();
                Uniprot uniEntry = (Uniprot)uniSpeciesMap.get (uniID);
                ArrayList <String> xrefList = new ArrayList <String> ();
                if (xrefType.equals (ExtResource.LIST_GI)) {
                    xrefList = uniEntry.getGIList();
                } else if (ExtResource.LIST_ID.equals (xrefType)) {
                    TreeMap <String, ExtResource> entrezMap =
                                     uniEntry.getXRef(ExtResource.RES_ENTREZ);
                    if (entrezMap != null) {
                        xrefList.addAll (entrezMap.keySet());
                    }
                }
                for (int j = 0; j < xrefList.size (); j++) {
                    String xrefID = xrefList.get (j);
                    String otherUniID = xrefUniSpeciesMap.get (xrefID);
                    if (otherUniID != null) {           //overwrite existing
                        xrefUniSpeciesMap.put (xrefID, otherUniID + GenUtil.SEMICOLON + uniID);
                    } else {
                        xrefUniSpeciesMap.put (xrefID, uniID);
                    }
                }
            }
            if (xrefType.equals (ExtResource.LIST_GI)) {
                GIUniMap.put(String.valueOf (speciesIDArr [i]), xrefUniSpeciesMap);
            } else if (xrefType.equals (ExtResource.LIST_ID)) {
                entrezUniMap.put (String.valueOf (speciesIDArr [i]), xrefUniSpeciesMap);
            }
        }
    }

    public static ArrayList <String> getUniprotIDByXref (String xrefType,
                                                         String xrefID) {

        return getUniprotIDByXref (xrefType, xrefID, -1, false);
    }

    public static ArrayList <String> getUniprotIDByXref (String xrefType,
                                     String xrefID, int speciesID) {

        return getUniprotIDByXref (xrefType, xrefID, speciesID, true);
    }

    //helper method
    private static ArrayList <String> getUniprotIDByXref (String xrefType,
                       String xrefID, int speciesID, boolean validateSpecies) {

        if (!ExtResource.LIST_GI.equals (xrefType) &&
            !ExtResource.LIST_ID.equals (xrefType)) {
            throw new IllegalArgumentException ("Unsupported x-ref type for Uniprot: "
                                                + xrefType);
        }
        GenUtil.validateString(xrefID);
        if (validateSpecies) {
            GenUtil.validatePositiveInt(speciesID);
        }
        ArrayList <String> uniIDList = new ArrayList <String> ();
        boolean speciesMismatch = false;
        String uniIDStr = "";
        TreeMap <String, String> xrefUniMap = null;
        if (ExtResource.LIST_GI.equals (xrefType)) {
            xrefUniMap = GIUniMap.get (String.valueOf (speciesID));
        } else if (ExtResource.LIST_ID.equals (xrefType)) {
            xrefUniMap = entrezUniMap.get (String.valueOf (speciesID));
        }
        if (xrefUniMap != null) {
            uniIDStr = xrefUniMap.get(xrefID);
            if (uniIDStr != null) {
                String uniIDArr [] = uniIDStr.split (GenUtil.SEMICOLON);
                for (int i = 0; i < uniIDArr.length; i++) {
                    uniIDList.add (uniIDArr [i]);
                }
            } else {
                speciesMismatch = true;
                uniIDStr = "";
            }
        }
        if (uniIDList.size () == 0) {  //either species mismatch or really not found
           TreeMap <String, ExtResource> uniprotMap = null;
           try {
               uniprotMap = Uniprot.mapToXref(xrefType, xrefID);
           } catch (Exception e) {
               e.printStackTrace();
               System.out.println ("Unable to process request for: " + xrefID);
           }
           Iterator iterator = uniprotMap.keySet().iterator();
           while (iterator.hasNext()) {
               Uniprot uniEntry = (Uniprot)uniprotMap.get((String)iterator.next());
               uniIDList.add (uniEntry.getID());
               uniIDStr += uniEntry.getID () + GenUtil.SEMICOLON;
           }
        }
        if (ExtResource.LIST_GI.equals (xrefType)) {
            lookupSupMappings (xrefType, xrefID);
            if (uniIDList.size () == 0) {
                //BatchEntrezReader.main(null);
                if (!unmatchedGIUniList.contains (xrefID)) {
                    unmatchedGIUniList.add (xrefID);
                }
            } else {
                if (!matchedGIUniList.contains (xrefID)) {
                    matchedGIUniList.add (xrefID);
                }
                if (speciesMismatch == true) {
                    //System.err.println ("Species mismatch: " + xrefType + " " +
                    //        xrefID + " " + uniIDStr + " " + speciesID);
                }
            }
        }
        if (ExtResource.LIST_ID.equals (xrefType)) {
            if (uniIDList.size () == 0) {
                if (!unmatchedEntrezUniList.contains (xrefID)) {
                    unmatchedEntrezUniList.add (xrefID);
                }
            } else {
                if (!matchedEntrezUniList.contains (xrefID)) {
                    matchedEntrezUniList.add (xrefID);
                }
                if (speciesMismatch == true) {
                    //System.err.println ("Species mismatch: " + xrefType + " " +
                    //        xrefID + " " + uniIDStr + " " + speciesID);
                }
            }
        }
        
        return uniIDList;
    }

    public static void dumpUnmatchedUni (String xrefType) {

        if (ExtResource.LIST_GI.equals (xrefType)) {
            System.out.println ("Matched GI in Uniprot: " + matchedGIUniList.size ());
            System.out.println ("Unmatched GI in Uniprot: " + unmatchedGIUniList.size());
            for (int i = 0; i < unmatchedGIUniList.size(); i++) {
                System.out.println ("No Uniprot match for GI: " + unmatchedGIUniList.get (i));
            }
        } else if (ExtResource.LIST_ID.equals (xrefType)) {
            System.out.println ("Matched Entrez in Uniprot: " + matchedEntrezUniList.size ());
            System.out.println ("Unmatched Entrez in Uniprot: " + unmatchedEntrezUniList.size());
            for (int i = 0; i < unmatchedEntrezUniList.size(); i++) {
                System.out.println ("No Uniprot match for Entrez: " + unmatchedEntrezUniList.get (i));
            }
        }
    }

    private static void lookupSupMappings (String xrefType, String xrefID) {

        if (!ExtResource.LIST_GI.equals (xrefType)) {
            return;
        }
        if (IdentifierMapperService.needSuppMappings) {
            IdentifierMapperService.needSuppMappings = false;
            loadSupMappings();
        }
    }

    private static void loadSupMappings () {

        BatchEntrezReader.clear();
        BatchEntrezReader.main (null);
        int cntRefSeq = 0, cnt = 0;
        try {
            //load protein mappings.
            TreeMap <String, String> IDProMap =
                       BatchEntrezReader.getMappedIDs(BatchEntrezReader.LOOKUP_PRO);
            Iterator iterator = IDProMap.keySet().iterator();
            while (iterator.hasNext()) {
                String key = (String)iterator.next();
                boolean isRefSeqPro = false;
                for (int i = 0; i < REFSEQ_PRO_PREFIX_ARR.length; i++) {
                    if (IDProMap.get (key).startsWith(REFSEQ_PRO_PREFIX_ARR[i])) {
                        isRefSeqPro = true;
                        break;
                    }
                }
                if (isRefSeqPro) {
                    IdentifierMapperService.GIRefSeqProMap.put (key, IDProMap.get (key));
                    cntRefSeq++;
                } else {
                    IdentifierMapperService.GIProMap.put (key, IDProMap.get (key));
                }
                cnt++;
            }
            System.out.println ("Count Pro Refseq: " + cntRefSeq + " " + cnt);
            cnt = 0;
            cntRefSeq = 0;

            //load output from RefSeq protein against Uniprot
            GIUniSupMap = BatchEntrezReader.getSupMappedIDs ();

            //load refseq mappings
            TreeMap <String, String> IDNucMap = BatchEntrezReader.getMappedIDs(
                                                BatchEntrezReader.LOOKUP_NUC);
            iterator = IDNucMap.keySet().iterator();
            while (iterator.hasNext()) {
                String key = (String)iterator.next();
                boolean isRefSeqMrna = false;
                for (int i = 0; i < REFSEQ_RNA_PREFIX_ARR.length; i++) {
                    if (IDNucMap.get (key).startsWith(REFSEQ_RNA_PREFIX_ARR[i])) {
                        isRefSeqMrna = true;
                        break;
                    }
                }
                if (isRefSeqMrna) {
                    IdentifierMapperService.GIRefSeqMrnaMap.put (key, IDNucMap.get (key));
                    cntRefSeq++;
                }
                cnt++;
            }
            System.out.println ("Count mRNA Refseq: " + cntRefSeq + " " + cnt);
         } catch (Exception e) {
            e.printStackTrace();
            System.err.println ("Unable to process request.");
        }
    }

    //using an arraylist for potential future refactoring.
    public static ArrayList <String> getXrefLO (String sourceType, String
                  sourceID, String targetType) {

        ArrayList <String> loList = new ArrayList <String> ();
        if (ExtResource.LIST_GI.equals (sourceType)) {
            if (ExtResource.LIST_REFSEQ_PRO.equals (targetType)) {
                if (IdentifierMapperService.GIRefSeqProMap.containsKey(sourceID)) {
                    loList.add (IdentifierMapperService.GIRefSeqProMap.get (sourceID));
                }
            } else if (ExtResource.LIST_REFSEQ_RNA.equals(targetType)) {
                if (IdentifierMapperService.GIRefSeqMrnaMap.containsKey(sourceID)) {
                    loList.add (IdentifierMapperService.GIRefSeqMrnaMap.get (sourceID));
                }
            } else if (IdentifierMapperService.LO_PROTEIN.equals (targetType)) {
                 if (IdentifierMapperService.GIProMap.containsKey(sourceID)) {
                    loList.add (IdentifierMapperService.GIProMap.get (sourceID));
                 }
            } else if (IdentifierMapperService.LO_GI_UNI.equals (targetType)) {
                if (IdentifierMapperService.GIUniSupMap.containsKey(sourceID)) {
                    loList.addAll (IdentifierMapperService.GIUniSupMap.get (sourceID));
                }
            }
        } else {
            throw new IllegalArgumentException ("Unsupported Left Over type: " +
                                                sourceType);
        }
        
        return loList;
    }
}
