package fmph.features.mitochondrion1;

import calhoun.analysis.crf.AbstractFeatureManager;
import calhoun.analysis.crf.CacheStrategySpec;
import calhoun.analysis.crf.CacheStrategySpec.CacheStrategy;
import calhoun.analysis.crf.FeatureList;
import calhoun.analysis.crf.FeatureManagerEdgeExplicitLength;
import calhoun.analysis.crf.FeatureManagerNodeBoundaries;
import calhoun.analysis.crf.ModelManager;
import calhoun.analysis.crf.features.supporting.phylogenetic.EvolutionaryModel;
import calhoun.analysis.crf.features.supporting.phylogenetic.HKY85Model;
import calhoun.analysis.crf.features.supporting.phylogenetic.Kimura80Model;
import calhoun.analysis.crf.features.supporting.phylogenetic.PhylogeneticTreeFelsensteinOrder;
import calhoun.analysis.crf.io.InputSequence;
import calhoun.analysis.crf.io.MultipleAlignmentInputSequence;
import calhoun.analysis.crf.io.MultipleAlignmentInputSequence.MultipleAlignmentColumn;
import calhoun.analysis.crf.io.TrainingSequence;

import calhoun.util.Assert;

import fmph.features.supporting.phylogenetic.CodonEvolutionaryModel;

import fmph.features.supporting.phylogenetic.GoldmanYang94;

import java.util.List;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

public class CodonLogprobMitochondrion1 extends AbstractFeatureManager<MultipleAlignmentColumn> implements FeatureManagerNodeBoundaries<MultipleAlignmentColumn> {

    private static final Log log = LogFactory.getLog(CodonLogprobMitochondrion1.class);


    int startIx; // The index of the first feature managed by this FeatureManager
    ModelManager model;

    EvolutionaryModel emodelNoncoding; // one model for intronic/intergenic regions
    CodonEvolutionaryModel emodelExonic;
    
    double codonEdgeCoef = 1;
    
    
    double[] emodelExonicParams = { 
            //codon freq
            0.06870156,0.01163551,0.06247777,0.00667686,
            0.03160051,0.00535196,0.02873777,0.00307114,
            0.03577897,0.00605133,0.00001000,0.00001000,
            0.02109063,0.00357198,0.01918000,0.00204972,
            0.02191554,0.00371168,0.01993017,0.00212989,
            0.01008044,0.00170726,0.00916724,0.00097968,
            0.01141654,0.00193354,0.01038230,0.00110953,
            0.00672783,0.00113945,0.00611835,0.00065385,
            0.06803812,0.01152315,0.06187443,0.00661238,
            0.03129535,0.00530028,0.02846025,0.00304148,
            0.03544335,0.00600280,0.03223248,0.00344461,
            0.02088696,0.00353748,0.01899478,0.00202993,
            0.04945072,0.00837513,0.04497090,0.00480594,
            0.02274574,0.00385229,0.02068517,0.00221058,
            0.02576055,0.00436289,0.02342686,0.00250358,
            0.01518083,0.00257107,0.01380558,//0.00147537,
            0.65390,//kappa
            0.08634,//omega
            4 //icode
            };
    
    
    
    double[] emodelNoncodingParams = { 
            
            0.43914528597477975,0.08378919793512597,0.1197832395921617,0.35728227649793254, //nucl. freq
            7.16, //transitions
            3.97 //transversions
            };

    //static KmerHasher hforward = new KmerHasher(KmerHasher.ACGTother, 1); // a character hasher for forward strand
    //static KmerHasher hbackward = new KmerHasher(KmerHasher.ACGTotherRC, 1); // a character hasher for reverse strand

    ///////////////////////////////////////////////////////////////////////////////

    public CodonLogprobMitochondrion1() {
    } // a constructor with no arguments

    public int getNumFeatures() { // there is exactly one feature
        return 2;
    }

    public String getFeatureName(int featureIndex) {
        
            String[] vals = new String[] { "codon logprob", "noncoding nucle logprob"};
            int feat = featureIndex - startIx;
            String table = vals[feat];
            return table + " phylogeny";
        
    }
    


    public void evaluateNode(InputSequence<? extends MultipleAlignmentColumn> seq, int pos, int state, FeatureList result) {
        if(pos < 2 || pos > seq.length()-3) return;
        Assert.a(state < model.getNumStates());

        double val = 0.0;
        int ephase;
        int featureOffset = Integer.MIN_VALUE;
        switch (state) {
        case 0:
            val = emodelNoncoding.logprob(seq.getX(pos), true);
            featureOffset = 0;
            break;
        case 1:
        case 2:
        case 3:
            ephase = ((pos - state + 1) % 3 + 3) % 3; //((pos-(state-1))%3 +3)%3;
            if(ephase==2){
                val = emodelExonic.logprob(seq.getX(pos-2),seq.getX(pos-1),seq.getX(pos),true);
                featureOffset = 1;
            }
            break;
        case 4:
        case 5:
        case 6:
            val = emodelNoncoding.logprob(seq.getX(pos), true);
            featureOffset = 0;
            
            break;
        case 7:
        case 8:
        case 9:
            ephase = ((-pos + state + 1) % 3 + 3) % 3; // ((-pos+2+(state-7))%3 +3)%3;
            if(ephase==2){
                val = emodelExonic.logprobRC(seq.getX(pos),seq.getX(pos+1),seq.getX(pos+2),true);
                featureOffset = 1; // + ephase;
            }
            break;
        case 10:
        case 11:
        case 12:
            val = emodelNoncoding.logprobRC(seq.getX(pos), true);
            featureOffset = 0;
            break;
        default:
            Assert.a(false);
        }
        if(featureOffset != Integer.MIN_VALUE)
            result.addFeature(startIx + featureOffset , val);
    }


    public void train(int startingIndex, ModelManager modelInfo, final List<? extends TrainingSequence<? extends MultipleAlignmentColumn>> data) {
        startIx = startingIndex;
        model = modelInfo;

        final PhylogeneticTreeFelsensteinOrder felsOrder = ((MultipleAlignmentInputSequence.MultipleAlignmentColumn)data.get(0).getX(0)).getMultipleAlignment().getFelsensteinOrder();
        double[] codonExonicPi = new double[64];
        System.arraycopy(emodelExonicParams,0,codonExonicPi,0,63);
        
        double sum = 0;
        for(int i=0;i<63;i++){
            sum += emodelExonicParams[i];
        }
        codonExonicPi[63] = 1 - sum;
        
        emodelNoncoding = 
            new EvolutionaryModel(
                felsOrder, new double[]{emodelNoncodingParams[0],emodelNoncodingParams[1],emodelNoncodingParams[2],emodelNoncodingParams[3]}, 
                new HKY85Model(new double[]{emodelNoncodingParams[4],emodelNoncodingParams[5], emodelNoncodingParams[0],emodelNoncodingParams[1],emodelNoncodingParams[2]}));
        log.debug("Trained conon noncoding feauture");
        emodelNoncoding.summarize();
        emodelExonic = new CodonEvolutionaryModel(felsOrder, codonExonicPi, new GoldmanYang94(emodelExonicParams), codonEdgeCoef);
        log.debug("Trained conon exonic feauture");
        emodelExonic.summarize();
    }


    @Override
    public CacheStrategySpec getCacheStrategy() {
            
            CacheStrategySpec css = new CacheStrategySpec(CacheStrategy.DENSE_NODE_BOUNDARY);
            CacheStrategySpec.DenseBoundaryCachingDetails details = new CacheStrategySpec.DenseBoundaryCachingDetails(13); // we will use 9 tables
            

            // If you want to predict for real and mesh with a PWM, then set pads, for example, as below:
            //      donor[j]    = new PWMLookup(3,6,pseudoCounts);   // donor signal           xxx|GTxxxx 
            //      acceptor[j] = new PWMLookup(9,6,pseudoCounts);   // acceptor signal  xxxxxxxAG|xxxxxx
            //  start = new PWMLookup(9,6,pseudoCounts);             // start signal xxxxxxxxx|ATGxxx
            //  stop  = new PWMLookup(3,9,pseudoCounts);             // stop signal        xxx|TAGxxxxxx
            
            // IDEALLY HAVE SOME ASSERTIONS TYING THESE PADS TO THE PWM SIZES
            details.add(0,0,startIx, 0, 0);
            details.add(1,1,startIx+1, 0, 2);
            details.add(2,2,startIx+1, 0, 2);
            details.add(3,3,startIx+1, 0, 2);
            details.add(4,4,startIx, 0, 0);
            details.add(5,5,startIx, 0, 0);
            details.add(6,6,startIx, 0, 0);
            
            details.add(7,7,startIx+1, 2, 0);  
            details.add(8,8,startIx+1, 2, 0);
            details.add(9,9,startIx+1, 2, 0);
            details.add(10,10,startIx, 0, 0);  
            details.add(11,11,startIx, 0, 0);
            details.add(12,12,startIx, 0, 0);
            
            
              
            // NOTE: The minimum lengths of the Semi-Markov training and inference for each state must be AT least as big as the sum of the two pads at either end.
            
            details.check();
                            
            css.details = details;
            
            return css;
    }


    public void setEmodelNoncodingParams(double[] emodelNoncodingParams) {
        this.emodelNoncodingParams = emodelNoncodingParams;
    }

    public double[] getEmodelNoncodingParams() {
        return emodelNoncodingParams;
    }

    public void setEmodelExonicParams(double[] emodelExonicParams) {
        this.emodelExonicParams = emodelExonicParams;
    }

    public double[] getEmodelExonicParams() {
        return emodelExonicParams;
    }

    public void setCodonEdgeCoef(double codonEdgeCoef) {
        this.codonEdgeCoef = codonEdgeCoef;
    }

    public double getCodonEdgeCoef() {
        return codonEdgeCoef;
    }
}
