import java.io.*;
import java.util.*;
import java.io.File;
import weka.core.*;
import weka.classifiers.*;
import weka.core.converters.ConverterUtils.DataSource;

public class PredictionModels {
    String prediction_type = null; // ALL, BACKBONE, SIDECHAIN

    private final static String MODEL_FILE_ALLATOMS       = "predmodels_all.sobj";
    private final static String MODEL_FILE_BACKBONE_ONLY  = "predmodels_backbone.sobj";
    private final static String MODEL_FILE_SIDECHAIN_ONLY = "predmodels_sidechain.sobj";

    private final static String OBJECT_MODEL_DIRECTORY = "/lib";
    private final static String MODEL_DIRECTORY = "/lib/predmodels";

    private static String SHIFTXP_DIR = ShiftXp.SHIFTX2_DIR;

    HashMap<String, Classifier> hsmap_predmodel = new HashMap<String, Classifier>();

    public Instances instanceHeader = null;

    PredictionModels() { 
       // setHeader();
    }

    PredictionModels(String type) { 
        this.prediction_type = type; 
       // setHeader();
    }

    // protected void setHeader() {
    public void setHeader() {

       // System.out.println("SHIFTXP_DIR:" + SHIFTXP_DIR+"\n\n");
       // System.out.println("SHIFTX2_DIR:" + ShiftXp.SHIFTX2_DIR+"\n\n");

        String headerfilename = ShiftXp.SHIFTX2_DIR + "/lib/data-header.arff"; 

        try {
            // assigned instance header from data-header.arff
            DataSource source = new DataSource( headerfilename );
            this.instanceHeader = source.getDataSet(); 
            this.instanceHeader.setClassIndex(this.instanceHeader.numAttributes()-1);
        } catch ( Exception e ) {
            System.err.println("\n********************************************************************");
            System.err.println("[Error]\n\t Cannot read header file(arff) for dataset (" + headerfilename + ")");
            System.err.println("********************************************************************\n\n");
            e.printStackTrace(); 
        }
    }
    
    public Instances getInstanceHeader() {
    	return this.instanceHeader;    	
    }

    public Classifier getModel(String atomName) {
        return hsmap_predmodel.get(atomName);
    }

    @SuppressWarnings({"unchecked"})
    public void readHashObjectPredModelFile(String fname) {
        try {
            System.out.println("Loading prediction models .... ");

            ObjectInputStream ohs = new ObjectInputStream( new BufferedInputStream( new FileInputStream( fname ) ));
            hsmap_predmodel = (HashMap<String,Classifier>) ohs.readObject();
            ohs.close();

        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        }
    }

    public void saveHashObjectPredModelFile(String fname) {
        try {
            System.out.println("Creating an object file for prediction models ...... ");

            ObjectOutputStream ohs = new ObjectOutputStream(new BufferedOutputStream( new FileOutputStream( fname ) ));
            ohs.writeObject(hsmap_predmodel);
            ohs.close();

        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        }
    }

    public void setHashMap() {
        Classifier classifier;

        String[] atomNameArray = { 
            "CA","CB","C","HA","H","N"
            ,"CD","CD1","CD2","CE","CE1","CE2","CG","CG1","CG2","CZ"
            ,"HA2","HA3","HB","HB2","HB3","HD1","HD2","HD3"
            ,"HE","HE1","HE2","HE3","HG","HG1","HG12","HG13","HG2","HG3","HZ"
            ,"HD21","HD22","HE21","HE22"  
        }; // no of models: 45  // CE3, CH2, CZ2, CZ3, HH2, HZ2, HZ3  eliminated

        int BACKBONE_MODEL_LENGTH = 6; 
        int SIDECHAIN_MODEL_LENGTH = 33;

        String[] modelFileArray = {
            "Dec20-AR-A20B20M3-Sp7N3-d20111018-refdbphycs2-CA.model",
            "Dec20-AR-A20B20M3-Sp6N3-d20111018-refdbphycs2-CB.model",
            "Dec23-AR-A20B20M3-Sp6N3-refdbphycs2-d20111223-C.model",
            "Dec20-AR-A20B20M3-Sp7N3-d20111018-refdbphycs2-HA.model",
            "Dec20-AR-A20B20M3-Sp8N3-d20111018-refdbphycs2-H.model",
            "Dec23-AR-A20B30M4-Sp5N3-d20111018-refdbphycs2-N.model",
            "AR-A10B25M3-d20111201-refdbcs-CD.model",
            "AR-A20B30M3-d20111205-paper-refdbcs-CD1.model",
            "AR-A15B25M3-d20111114-refdbcs-CD2.model",
            "AR-A10B25M3-d20111201-refdbcs-CE.model",
            "AR-B20M3-d20111114-refdbcs-CE1.model",
            "AR-A10B20M3-d20111201-refdbcs-CE2.model",
            "AR-A10B25M3-d20111201-refdbcs-CG.model",
            "AR-A10B25M3-d20111201-refdbcs-CG1.model",
            "AR-A15B25M3-d20111201-refdbcs-CG2.model",
            "AR-A10B20M3-d20111201-refdbcs-CZ.model",
            "AR-B20M3-d20111114-refdbcs-HA2.model",
            "AR-A10B20M3-d20111201-refdbcs-HA3.model",
            "AR-A10B20M3-d20111201-refdbcs-HB.model",
            "AR-A15B20M3-d20111114-refdbcs-HB2.model",
            "AR-A20B25M3-d20111114-refdbcs-HB3.model",
            "AR-B20M3-d20111114-refdbcs-HD1.model",
            "AR-B20M3-d20111114-refdbcs-HD2.model",
            "AR-A10B25M3-d20111201-refdbcs-HD3.model",
            "AR-A15B25M3-d20111201-refdbcs-HE.model",
            "AR-A10B25M3-d20111201-refdbcs-HE1.model",
            "AR-A10B20M3-d20111201-refdbcs-HE2.model",
            "AR-A10B20M3-d20111201-refdbcs-HE3.model",
            "AR-A15B25M3-d20111201-refdbcs-HG.model",
            "AR-A20B25M3-d20111114-refdbcs-HG1.model",
            "BR-B30M2-d20111201-refdbcs-HG12.model",
            "AR-A20B40M3-d20111206-refdbcs-NoLimit-HG13.model",
            "AR-A10B20M2-d20111114-refdbcs-HG2.model",
            "AR-A15B25M3-d20111114-refdbcs-HG3.model",
            "AR-A20B25M3-d20111114-refdbcs-HZ.model",
            "AR-B20M3-d20111114-refdbcs-HD21.model",
            "AR-A10B20M3-d20111201-refdbcs-HD22.model",
            "AR-A15B25M3-d20111201-refdbcs-HE21.model",
            "AR-A10B25M3-d20111201-refdbcs-HE22.model",
        };

        try {
            ObjectInputStream ois = null;

            // if option is not sidechain only, put backbone models into modeling file 
            if ( ! this.prediction_type.equals("SIDECHAIN") ) {
                for(int i=0; i < BACKBONE_MODEL_LENGTH; i++)  {
                    ois = new ObjectInputStream( new BufferedInputStream( new FileInputStream( ShiftXp.SHIFTX2_DIR + MODEL_DIRECTORY + "/" + modelFileArray[i]) ));
                    classifier = (Classifier) ois.readObject();
                    hsmap_predmodel.put(atomNameArray[i],classifier);
                    ois.close();
                }
            }

            // if option is not backbone only, put/append sidechain models into modeling file
            if ( ! this.prediction_type.equals("BACKBONE") ) {
                for(int i=BACKBONE_MODEL_LENGTH; i < (BACKBONE_MODEL_LENGTH + SIDECHAIN_MODEL_LENGTH); i++)  {
                    ois = new ObjectInputStream( new BufferedInputStream( new FileInputStream( ShiftXp.SHIFTX2_DIR + MODEL_DIRECTORY + "/" + modelFileArray[i]) ));
                    classifier = (Classifier) ois.readObject();
                    hsmap_predmodel.put(atomNameArray[i], classifier);
                    ois.close();
                }
            }

        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        }
    }

    public void loadPredictionModels(String type) {
        this.prediction_type = type;

        String model_file = null;

        setHeader();   // initialize weka instance's header using arff file
 
        /**         
            if the loading time is not different, it will be used just one full mode file
        */

        if ( this.prediction_type.equals("ALL") == true) {
            model_file = ShiftXp.SHIFTX2_DIR + OBJECT_MODEL_DIRECTORY + "/" + MODEL_FILE_ALLATOMS;
        } else if ( this.prediction_type.equals("BACKBONE") == true) {
            model_file = ShiftXp.SHIFTX2_DIR + OBJECT_MODEL_DIRECTORY + "/" + MODEL_FILE_BACKBONE_ONLY;
        } else if ( this.prediction_type.equals("SIDECHAIN") == true) {
            model_file = ShiftXp.SHIFTX2_DIR + OBJECT_MODEL_DIRECTORY + "/" + MODEL_FILE_SIDECHAIN_ONLY;
        } 

        File f = new File( model_file );

        if ( f.exists() ) {
            readHashObjectPredModelFile( model_file );
        } else {
            setHashMap(); 
            saveHashObjectPredModelFile( model_file ); 
        }
    }

}

