package org.baderlab.lola.view;

import org.baderlab.lola.model.ProfileSet;
import org.baderlab.lola.model.LogoTree;
import org.baderlab.lola.model.LogoGenerator;
import org.baderlab.lola.model.Profile;
import org.baderlab.lola.LolaParameterSet;
import org.baderlab.brain.*;

import cytoscape.task.ui.JTaskConfig;
import cytoscape.task.Task;
import cytoscape.task.TaskMonitor;


import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.awt.image.BufferedImage;

/**
 Copyright 2007 Bader Lab

 This program is free software; you can redistribute it and/or modify
 it under the terms of the GNU Lesser General Public License as published by
 the Free Software Foundation; either version 3 of the License, or
 (at your option) any later version.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program.  If not, see <http://www.gnu.org/licenses/>.

 */

/**
 * Created by IntelliJ IDEA.
 * User: moyez
 * Date: May 28, 2007
 * Time: 6:14:41 PM
 * To change this template use File | Settings | File Templates.
 */
public class GenerateLogoTask implements Task {
    private TaskMonitor taskMonitor = null;
    BrainAlgorithm alg = null;

    ProfileSet profileSet;
    LolaParameterSet paramSet;

    public GenerateLogoTask(ProfileSet profileSet, LolaParameterSet paramSet) {
        this.profileSet = profileSet;
        this.paramSet = paramSet;
        alg = new BrainAlgorithm();
    }

    public void run() {
        if (taskMonitor == null) {
            throw new IllegalStateException("Task Monitor is not set.");
        }
        alg.setTaskMonitor(taskMonitor);
        taskMonitor.setStatus("Loading profiles");

        //determine if we're dealing with a project file or a single protein profile
        //read profile(s) from the profile file using given parameters
        // (a single protein profile or a set of profiles linked from a project file)
        List profileList = null;
        try
        {
            profileList = PeptideToProfileReader.readPeptidesAsProfiles(
                paramSet.getProfileFile(),
                paramSet.getFuzzFactor(),
                paramSet.getCodonBiasFile(),
                paramSet.getUniquePeptides(), true);
        }
        catch(Exception e)
        {
            // try reading in as a PWM file
            profileList = PWMToProfileReader.readPWMsAsProfiles(
                paramSet.getProfileFile(),
                //0.01,
                paramSet.getCodonBiasFile());
        }
        //generate logo image(s) for the profile(s)
        profileSet.clear(); // only one project file can be processed at a time, so clear the profiles first
        taskMonitor.setStatus("Drawing logos");
        for (int i = 0; i < profileList.size(); i++) {
            ProteinProfile proteinProfile = (ProteinProfile) profileList.get(i);

            //generate logo image
            BufferedImage image = LogoGenerator.generateImage(
                    proteinProfile,
                    paramSet,
                    "_logo");

            //set the logo profile and add it to the profile set
            Profile profile = new Profile();
            profile.set(proteinProfile);
            profile.setImage(image);
            profileSet.addElement(profile);

        }    }

    /**
     * Saves a 20aa biojava alphabet for ease of use (ignore the U symbol)
     */
    /*
    @SuppressWarnings("unchecked")
	private ArrayList get20aaAlphabet() {
        ArrayList alphabet = new ArrayList(20);
        HashMap alphabetMap = ProteinSequenceUtil.get20aaAlphabet();
        Collection symbols = alphabetMap.values();
        for (Iterator iterator = symbols.iterator(); iterator.hasNext();) {
            Symbol symbol = (Symbol) iterator.next();
            alphabet.add(symbol);
        }
        return alphabet;
    }
    */

    /**
     * Re-weight the profile to correct for the presence of 'X' ambiguity symbol. BioJava 1.4 interprets this symbol as
     * the simultaneous presence of all 20 amino acids, hence adding 1 count to each AA symbol in the distribution. 
     * Here we correct this by removing 1 count and adding 1/20th of a count (0.05) to each symbol, for each occurrence
     * of 'X' in the distribution of a given alignment position. (Dharsee 20090603) 
     */
    /*
    public void correctForAmbiguitySymbolX(ProteinProfile proteinProfile) {

    	//get the amino acid alphabet
    	ArrayList alphabet = get20aaAlphabet();
        if (alphabet == null) {
            throw new IllegalStateException("Alphabet array could not be set.");
        }
        WeightMatrix wm = proteinProfile.getWeightMatrix();
        Alignment align = new SimpleAlignment(proteinProfile.getSequenceHashMap());

        //calculate the sum of weights of the unnormalised distributions
        Distribution[] dists = null;
        ArrayList<Double> distWeightSum = new ArrayList<Double>();		//array list to hold the sum of weights at each position
        try {
			dists = DistributionTools.distOverAlignment(align, false, paramSet.getFuzzFactor());
		} catch (IllegalAlphabetException e1) {
			e1.printStackTrace();
		}
		for (int i = 0; i < dists.length; i++) {
			double weightSum = 0;
        	Distribution d = dists[i];
            Iterator l = alphabet.iterator();
            while (l.hasNext()) {
                Symbol symbol = (Symbol) l.next();
                try {
					double weight = d.getWeight(symbol);
					weightSum += weight;
				} catch (IllegalSymbolException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
            }
            distWeightSum.add(i, weightSum);
		}
        //calculate the number of 'X' symbols occurring at each alignment position
        List seqs = align.getLabels();
        ArrayList<Integer> xFrequency = new ArrayList<Integer>();	//array list to hold the X symbol count at each position
        for (int i = 0; i < align.length(); i++) {
        	xFrequency.add(i, 0);
        }
        for (int i = 0; i < align.length(); i++) { //for each alignment position
        	for(Iterator j = seqs.iterator(); j.hasNext();){// for each sequence
                Object seqLabel = j.next();
                Symbol s = align.symbolAt(seqLabel, i + 1);
                try {
					String symbolString = ProteinTools.getAlphabet().getTokenization("token").tokenizeSymbol(s);
	                if (symbolString.equalsIgnoreCase("X")) {
	                	xFrequency.set(i, xFrequency.get(i) + 1);                	
	                }
				} catch (IllegalSymbolException e) {
					e.printStackTrace();
				} catch (BioException e) {
					e.printStackTrace();
				}
        	}
        }
        
        //debug/testing
        System.out.println("WM Columns: " + wm.columns() + " | Alignment Length: " + align.length());
        for (int i = 0; i < xFrequency.size(); i++) {
        	System.out.println("Xs at position " + i + ": " + xFrequency.get(i) + " | weightSum = " + distWeightSum.get(i));
        }        
        DistributionTrainerContext dtc = new SimpleDistributionTrainerContext();
        try {
            for (int i = 0; i < wm.columns(); i++) {
                Distribution d = wm.getColumn(i);
                double weight = 0.0;
                Symbol symbolX = AlphabetManager.getAllAmbiguitySymbol(ProteinTools.getAlphabet());
                String symbolStringX = ProteinTools.getAlphabet().getTokenization("token").tokenizeSymbol(symbolX);
                dtc.registerDistribution(d);
            	//double weightX = dtc.getCount(d, symbolX);
                double weightX = d.getWeight(symbolX);
            	System.out.println("weight(" + symbolStringX + ")[" + i + "] = " + weightX);
                //divide all weights by the number of codons, keeping track of the sum of the new weights
                Iterator l = alphabet.iterator();
                while (l.hasNext()) {
                    Symbol symbol = (Symbol) l.next();
                    String symbolString = ProteinTools.getAlphabet().getTokenization("token").tokenizeSymbol(symbol);
                    if (!symbolString.equals("U")) {
                    	int xFreq = xFrequency.get(i);
                    	double weight0 = d.getWeight(symbol);
                        weight = d.getWeight(symbol) - xFreq + (xFreq/alphabet.size());
                        System.out.println("weight0(" + symbolString + ") = " + weight0 + " | " + "weight1(" + symbolString + ") = " + weight);
                        //d.setWeight(symbol, weight);
                    }
                }
            }
        } catch (BioException e) {
            e.printStackTrace();
//        } catch (ChangeVetoException e) {
//            e.printStackTrace();
        }
    }
    */
    
    public void halt() {
        //TODO - implement a cancel method in BrainAlgorithm that will halt logo tree drawing
    }

    public void setTaskMonitor(TaskMonitor taskMonitor) throws IllegalThreadStateException {
        if (this.taskMonitor != null) {
            throw new IllegalStateException("Task Monitor is already set.");
        }
        this.taskMonitor = taskMonitor;
    }

    public String getTitle() {
        return new String("Generating Sequence Logos");
    }

}
