import java.util.ArrayList;
import java.util.Collections;
import weka.core.*;
import weka.classifiers.*;

public class CSPredictionWekaModel {
    protected FeatureReader ds;
    protected RandomCoil rcoil;
    protected LimitedCShift chklimcs;
    protected String outputFormat;
    protected String atomtype;
    protected PredictionModels predmodels;
    
    private boolean DEBUG = ShiftXp.DEBUG;
    private boolean RCOIL_PRINTOUT = ShiftXp.RCOIL_PRINTOUT;

    public CSPredictionWekaModel(PredictionModels predmodels, FeatureReader data, RandomCoil rc, LimitedCShift limcs, String outform, String atomtype) {
        this.predmodels = predmodels;
        this.ds = data;
        this.rcoil = rc;
        this.chklimcs = limcs;
        this.outputFormat = outform;
        this.atomtype = atomtype;
    }


    /** 
      ? classifier: how to save/restore classifiers generated by weka
      vector : use index that matched atomname
      hashmap: use string atomname   
    */

    private double predcs_classifierModel(Residue res, Atom atom) {
        Classifier classifier;        
        double predcs = -9999;
        double predphycs = -9999; 

        FeatureToInstance f2ins = new FeatureToInstance();

        /**  set Instance using read features to fit weka model, norm2bin, normalization?  */
        Instance inst = f2ins.setInstanceWithTransform(res, atom);
        // OR
        // Instance inst = f2ins.setInstanceWithNormalizeTransform(res, atom);
        
       // inst.setDataset(ShiftXp.predmodels.instanceHeader);
        
        try {
                classifier = predmodels.getModel(atom.atomName);
                
                if ( classifier != null ) { 
                   // System.out.println("Called classifier for " + atom.atomName);
                        
                    predphycs = classifier.classifyInstance( inst );
                    
                    if( atom.atomName.equals("C") || atom.atomName.equals("CA") || atom.atomName.equals("CB") 
                       || atom.atomName.equals("H") || atom.atomName.equals("HA") || atom.atomName.equals("N") ) {
                        predcs = predphycs + (atom.rcoil + atom.rcurr + atom.electrostat + atom.hbondeffect);
                    } else { 
                        predcs = predphycs ;
                    }
                    
                    // System.out.println(">> predcs: " + predcs + "predphycs:"+ predphycs +" ("+ res.resNum + ":" + atom.atomName +")");
                } else {                    
                    // System.out.println("[Error] No classifier assigned for " + atom.atomName +"["+res.resNum+":"+res.resName+"]");
                    // System.exit(-1);
                }
            
        } catch (Exception e) {
            System.err.println("\n*******************************************************************");
            System.err.println("[Error] Exception in classify instance of prediction model");
            System.err.println("*******************************************************************\n\n");

            if (DEBUG) { e.printStackTrace(); }
            System.exit(-1);            
        }

        return predcs;
    } 

    public void dopredcs(String dirname, String pdbfname, Boolean deuteratedpdb) {
        Residue res;
        Atom atom;

        OutFile ofile = new OutFile();

        if ( ds.residues.isEmpty() )  {
            System.err.println("\n*******************************************************************");
            System.err.println("Error: No Residues in Input PDB");
            System.err.println("\n*******************************************************************");
            System.exit(-1);
        }

        StringBuffer outsb =  new StringBuffer();

        ArrayList<Integer> keys = new ArrayList<Integer>(ds.residues.keySet());
        Collections.sort(keys);

        // for tabular format print out of sidechain prediction
        StringBuffer outsbsideCs =  new StringBuffer();
        StringBuffer outsbsideHs =  new StringBuffer();
        double[] bbpcs   = new double[6];
        double[] sidepcs = new double[30];
        int atompos; 
        String phosphoResName=null, phospho_iupac=null;
        double phospho_side_rcoil, side_rcoil;

        for(Integer key : keys){ // Loop for Residue (key: residue number)
            res = ds.residues.get(key);
            ArrayList<String> atomNameList = new ArrayList<String>();

            // if Phosphorylated Residue
            if (res.resName.equals("SER") || res.resName.equals("THR") || res.resName.equals("TYR")) {
                phosphoResName = ds.hashmap_phosphoRes.get(res.resNum);
                // System.out.println("Phosphorylated:" + phosphoResName);
            } else {
                phosphoResName = null;
            }

            // Because of printing order of atoms
            if (phosphoResName == null) {
                atomNameList = rcoil.getAtoms(res.iupac, this.atomtype);
            } else {
                // set phosphorylated atome type and ipuac (pS, pT, pY)
                if (res.resName.equals("SER"))      phospho_iupac = "pS";
                else if (res.resName.equals("THR")) phospho_iupac = "pT";
                else if (res.resName.equals("TYR")) phospho_iupac = "pY";
                else phospho_iupac = null; 
                
                atomNameList = rcoil.getAtoms(phospho_iupac, this.atomtype);
            }
            // System.out.println("AtomNameList:");
            // System.out.println(atomNameList);
      

            // Side Chain Atoms ************************************************************
            if ( outputFormat.equals("TABULAR") ) { 
                for (int i=0; i < bbpcs.length; i++ )    bbpcs[i]   = -9999;
                for (int i=0; i < sidepcs.length; i++ )  sidepcs[i] = -9999;
            }

            Boolean haspredcs = true;


            // repeat for atoms in the current residue
            for(String k : atomNameList){
                atom = res.atoms.get(k);

                // *****************************************************************
                // check mismatched case
                if ( atom == null  && DEBUG) {
                    System.out.println("ERROR: mismatch atom name and residue type");
                    System.out.println("\tatom is null: " + res.resNum + "   res:"+res.resName+" ("+res.iupac+")   atom name:"+k);
                    System.out.println("\tatomnamelist: " + atomNameList);
                     continue;
                    // System.exit(-1);         
                } else if (atom == null & !DEBUG) {
                    System.out.println("ERROR: found mismatch atom name in residue type");
                    System.out.println("\tatom ("+k+") is null for" + res.resNum + " res:"+res.resName+" ("+res.iupac+")" );
                    System.out.println("\tatomnamelist: " + atomNameList);
                    System.exit(-1);            
                }

                // prediction chemical shift using models
                atom.predcs = predcs_classifierModel( res, atom );   // with weka prediction model by each atom
                // System.out.println(">> Atom:" + atom.atomName + " predcs:" + atom.predcs);
            
                // if Phosphorylated Residue, append atom P at the end
                if(phosphoResName != null && atom.atomName.equals("P")) {    
                    atom.predcs = ds.getPhosphoCS(phosphoResName, "P");
                    // System.out.println(">>  Phosphorylated predcs:" + atom.predcs);
                } else if(phosphoResName != null &&
                            ! ( atom.atomName.equals("C") || atom.atomName.equals("CA") || atom.atomName.equals("CB") 
                                || atom.atomName.equals("H") || atom.atomName.equals("HA") || atom.atomName.equals("N") )) {
                            // if Phosphorylated Residue, adjust sidechain atoms 
                    side_rcoil = ds.rcoil.getRCoil(res.iupac, atom.atomName);
                    phospho_side_rcoil = ds.getPhosphoCS(phosphoResName, atom.atomName);
                    atom.predcs = atom.predcs - side_rcoil + phospho_side_rcoil;
                    // System.out.println(">>  Phosphorylated side predcs:" + atom.predcs);
                } 
            
                // OPTION: deutured pdb checked
                if ( deuteratedpdb ) { 
                   atom.predcs = AdjustedDeuterated(atom.atomName, atom.predcs);
                }
               
                atom.predcs = chklimcs.checkLimitedCShiftBound(res.iupac, atom.atomName, atom.predcs);    

                if ( outputFormat.equals("CSV") ) { 

                    // print out random coil to test  (add 14 Sep 2011)
                    if (RCOIL_PRINTOUT == false) { 
                        if (phosphoResName == null) {
                            outsb.append(res.resNum +","+ res.resName + "," + res.iupac + ","  
                                + atom.atomName +","+ String.format("%.4f",atom.predcs) + "\n" );
                                //  "-"+res.ss + ","+res.hbond+ "\n" );
                        } else {
                            outsb.append(res.resNum +","+ phosphoResName + "," + phospho_iupac + ","  
                                + atom.atomName +","+ String.format("%.4f",atom.predcs) + "\n" );
                        }
                    } 
                    else {
                        if (phosphoResName == null) {
                            outsb.append(res.resNum +","+ res.resName + "," + res.iupac + ","  
                                + atom.atomName +","+ String.format("%.4f",atom.predcs) + "," + atom.rcoil + "\n" );
                                //  "-"+res.ss + ","+res.hbond+ "\n" );
                        } else {
                             outsb.append(res.resNum +","+ res.resName + "," + res.iupac + ","  
                                + atom.atomName +","+ String.format("%.4f",atom.predcs) + "," + atom.rcoil + "\n" );
                        }
                    }

                } else if ( outputFormat.equals("BMRB") ) { 

                        outsb.append(String.format("%4d",res.resNum) +" "+ res.resName + " " + res.iupac + " "  
                           + String.format("%-4s",atom.atomName) +" " + String.format("%8.4f",atom.predcs) 
                           + "\n" );

                } else {  
                    // tabular format
          
                    if ( ! ( atom.atomName.equals("HA2") || atom.atomName.equals("HA3") ) ) {
                        if ( "HA,H,N,CA,CB,C".contains(atom.atomName) ) {
                        //assign backbone atom array
                        atompos = getAtomPos(atom.atomName); 
                        if ( atompos >= 0 ) bbpcs[ atompos ] = atom.predcs;
                        } else  {
                        // assign sidechain atom array
                        atompos = getAtomPos(atom.atomName); 
                        if ( atompos >= 0 ) sidepcs[ atompos ] = atom.predcs;
                        }
                    }
                }
            } // end of residue for loop

            
            // to tabulat results : appending predicted chemical shifts for all atoms by residue 
            if ( outputFormat.equals("TABULAR") ) { 

                outsb.append(String.format("%4d",res.resNum) +"  "+ String.format("%-3s", res.iupac) );
                for ( int i=0; i < bbpcs.length; i ++ ) {
                    if ( bbpcs[i] != -9999 ) {   
                        outsb.append(" " + String.format("%.4f", bbpcs[i]));
                    } else {
                        if( i <= 1 ) 
                            outsb.append(" " + String.format("%6s", "----"));
                        else if ( i == 2 || i == 5 ) 
                            outsb.append(" " + String.format("%8s", "----"));
                        else 
                            outsb.append(" " + String.format("%7s", "----"));
                    }   
                }
                outsb.append("\n");

                outsbsideCs.append(String.format("%4d",res.resNum) +"  "+ String.format("%-3s", res.iupac) );
                outsbsideHs.append(String.format("%4d",res.resNum) +"  "+ String.format("%-3s", res.iupac) );
                // sidepcs.length - 3: HH2, HZ2, HZ3
                for ( int i=0; i < sidepcs.length-3; i++ ) {
                    if (i <= 5) {   
                        if ( sidepcs[i] != -9999 ) 
                             outsbsideCs.append(" " + String.format("%6.2f", sidepcs[i]));
                        else outsbsideCs.append(" " + String.format("%6s","---"));
                    } else {    
                        if ( sidepcs[i] != -9999 ) 
                             outsbsideHs.append(" " + String.format("%4.2f", sidepcs[i]));
                        else outsbsideHs.append(" " + String.format("%4s","---"));
                    }
                }
                outsbsideCs.append("\n");
                outsbsideHs.append("\n");
            }
        } // End of Loop Residue

        if ( outputFormat.equals("TABULAR") ) { 
            ofile.savePredCSToFileTabular(dirname + pdbfname + ".sxp", outsb, outsbsideCs, outsbsideHs);
        } else {
            ofile.savePredCSToFileBMRB(dirname + pdbfname + ".sxp", outsb);
        } 

    }

    private int getAtomPos(String atomname) {
        int pos;

        if (atomname.equals("HA"))       pos = 0;
        else if (atomname.equals("HA2")) pos = 0;
        else if (atomname.equals("HA3")) pos = 0;
        else if (atomname.equals("H"))   pos = 1;
        else if (atomname.equals("N"))   pos = 2;
        else if (atomname.equals("CA"))  pos = 3;
        else if (atomname.equals("CB"))  pos = 4;
        else if (atomname.equals("C"))   pos = 5;
        else if (atomname.equals("CD"))   pos = 0;
        else if (atomname.equals("CD1"))  pos = 1;
        else if (atomname.equals("CD2"))  pos = 2;
        else if (atomname.equals("CG"))   pos = 3;
        else if (atomname.equals("CG1"))  pos = 4;
        else if (atomname.equals("CG2"))  pos = 5;
        else if (atomname.equals("HB"))   pos = 6;
        else if (atomname.equals("HB2"))  pos = 7;
        else if (atomname.equals("HB3"))  pos = 8;
        else if (atomname.equals("HD1"))  pos = 9;
        else if (atomname.equals("HD2"))  pos = 10;
        else if (atomname.equals("HD21")) pos = 11;
        else if (atomname.equals("HD22")) pos = 12;
        else if (atomname.equals("HD3"))  pos = 13;
        else if (atomname.equals("HE"))   pos = 14;
        else if (atomname.equals("HE1"))  pos = 15;
        else if (atomname.equals("HE2"))  pos = 16;
        else if (atomname.equals("HE21")) pos = 17;
        else if (atomname.equals("HE22")) pos = 18;
        else if (atomname.equals("HE3"))  pos = 19;
        else if (atomname.equals("HG"))   pos = 20;
        else if (atomname.equals("HG1"))  pos = 21;
        else if (atomname.equals("HG12")) pos = 22;
        else if (atomname.equals("HG13")) pos = 23;
        else if (atomname.equals("HG2"))  pos = 24;
        else if (atomname.equals("HG3"))  pos = 25;
        else if (atomname.equals("HZ"))   pos = 26;
        else if (atomname.equals("HH2"))  pos = 27;
        else if (atomname.equals("HZ2"))  pos = 28;
        else if (atomname.equals("HZ3"))  pos = 29;
        else  pos = -9999;

        return pos;
    }   

    private double AdjustedDeuterated(String atomname, double predcs) {
        double adjustedpredcs = predcs;

        if ( atomname.equals("CA") ) {
             adjustedpredcs = predcs - 0.43;
        } else if ( atomname.equals("N") ) {
             adjustedpredcs = predcs - 0.23;
        } else if ( atomname.equals("CB") ) {
             adjustedpredcs = predcs - 0.82;
        } else {
             adjustedpredcs = predcs;
        }

        return adjustedpredcs;
    }
    
}
