/*
 * Created on May 19, 2005
 */
package calhoun.gebo.internal.db;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.Collection;
import java.util.Date;
import java.util.Iterator;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.hibernate.FlushMode;
import org.hibernate.Hibernate;

import calhoun.bean.AnalysisRun;
import calhoun.bean.AnnotatedTranscript;
import calhoun.bean.BioSequence;
import calhoun.bean.CurationFlag;
import calhoun.bean.Exon;
import calhoun.bean.Feature;
import calhoun.bean.Gene;
import calhoun.bean.GenomeAssembly;
import calhoun.bean.RetiredFeature;
import calhoun.bean.RetirementCode;
import calhoun.bean.Transcript;
import calhoun.bean.TranscriptType;
import calhoun.bean.UtrElement;
import calhoun.service.AnalysisRunService;
import calhoun.util.Assert;
import calhoun.util.DbSession;
import calhoun.util.ErrorException;

/**
 * Encapsulates database access used by Argo for creating/updating/retiring
 * transcripts. Interface is pretty ugly, but it's intended to just be a drop-in
 * replacement for the annotation PL/SQL package.
 * 
 * @author Michael Koehrsen
 */
public class ArgoTranscriptService {

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

    public static final int INSERT_MODE = 1;
    public static final int UPDATE_MODE = 2;

    private static final String[] EMPTY_STRINGS = {};

    private Connection connection;

    public ArgoTranscriptService(Connection connection) {
        this.connection = connection;
    }

    /**
     * Make sure we can validate the inserted/updated transcript
     */
    public void validateTranscriptCreation(int mode, String sequenceId,
            int start, int stop, String strand, int startCodon, int stopCodon,
            int transcriptTypeId, String transcriptName,
            String primaryEvidenceFeatureId, String transcriptComments,
            String createdBy, String modifiedBy, String[] curationFlagIds,
            String[] exonList, String[] supplementalEvidenceIds, String symbol,
            String geneId, String transcriptId) throws SQLException {
        DbSession sess = new DbSession(connection);
        // sess.setFlushMode(FlushMode.NEVER);
        try {
            Transcript transcript = createTranscript(sess, mode, sequenceId,
                    start, stop, strand, startCodon, stopCodon,
                    transcriptTypeId, transcriptName, primaryEvidenceFeatureId,
                    transcriptComments, createdBy, modifiedBy, curationFlagIds,
                    exonList, supplementalEvidenceIds, symbol, geneId,
                    transcriptId);
            // sess.rollback();
            // sess.close();
            sess.evict(transcript.getGene());
            sess.evict(transcript);
        } catch (Throwable t) {
            log.error(t);
            if (t instanceof RuntimeException)
                throw (RuntimeException) t;
            else
                throw new ErrorException(t);
        } finally {
            System.out.println("Please transcript go away!");
            // sess.flush();
            // sess.release();
            // sess.reset();
            sess.rollback();
            sess.close();
        }
    }

    /**
     * Actually insert/update the transcript.
     */
    public String[] insertOrUpdateAnnotatedTranscript(int mode,
            String sequenceId, int start, int stop, String strand,
            int startCodon, int stopCodon, int transcriptTypeId,
            String transcriptName, String primaryEvidenceFeatureId,
            String transcriptComments, String createdBy, String modifiedBy,
            String[] curationFlagIds, String[] exonList,
            String[] supplementalEvidenceIds, String symbol, String geneId,
            String transcriptId) throws SQLException {
        DbSession sess = new DbSession(connection);
        sess.setFlushMode(FlushMode.COMMIT);
        Transcript transcript = createTranscript(sess, mode, sequenceId, start,
                stop, strand, startCodon, stopCodon, transcriptTypeId,
                transcriptName, primaryEvidenceFeatureId, transcriptComments,
                createdBy, modifiedBy, curationFlagIds, exonList,
                supplementalEvidenceIds, symbol, geneId, transcriptId);
        saveTranscript(sess, transcript);
        return new String[] { transcript.getGene().getId(), transcript.getId() };
    }

    public void retireAnnotatedTranscript(String geneId, String transcriptId,
            String sequenceId, String strand, String annotator,
            String comments, int retirementCode) throws SQLException {

        DbSession sess = new DbSession(connection);
        try {
            BioSequence seq = (BioSequence) sess.getById(BioSequence.class,
                    sequenceId, false);
            Assert.a(seq != null, "No sequence found: " + sequenceId);

            AnalysisRunService ars = new AnalysisRunService(sess);
            AnalysisRun run = ars.fetchOrCreateAnalysisRun(seq.getGroup()
                    .getId(), "retire");

            Gene gene = (Gene) sess.getById(Gene.class, geneId, false);
            Assert.a(gene != null, "No gene found: " + geneId);

            AnnotatedTranscript transcript = (AnnotatedTranscript) sess
                    .getById(AnnotatedTranscript.class, transcriptId, false);

            transcript.populateCurationFlags(sess);

            Assert
                    .a(transcript != null, "No transcript found: "
                            + transcriptId);

            Date createDate = transcript.getCreationDate();
            Date modDate = transcript.getModificationDate();
            if (modDate == null)
                modDate = new Date();
            // Transcript is "new" if less than two hours old:
            boolean newTrans = (modDate.getTime() - createDate.getTime() < (2 * 60 * 60 * 1000));

            if (newTrans) {
                gene.removeTranscript(transcript);
                sess.delete(transcript);
                if (gene.getTranscripts().size() == 0) {
                    sess.delete(gene);
                }
            } else {
                sess.save(createRetiredFeature(sess, transcript, annotator,
                        comments, retirementCode));
                transcript.setSubtype("RETIRED_ANNOTATED_TRANSCRIPT");
                transcript.setAnalysisRun(run);
                Iterator iter = transcript.getExons().iterator();
                while (iter.hasNext()) {
                    Exon exon = (Exon) iter.next();
                    exon.setSubtype("RETIRED_ANNOTATED_EXON");
                    exon.setAnalysisRun(run);
                }

                boolean allTranscriptsRetired = true;
                iter = gene.getTranscripts().iterator();
                while (iter.hasNext()) {
                    if (!((Transcript) iter.next()).getSubtype().equals(
                            "RETIRED_ANNOTATED_TRANSCRIPT")) {
                        allTranscriptsRetired = false;
                    }
                }

                if (allTranscriptsRetired) {
                    sess.save(createRetiredFeature(sess, gene, annotator,
                            comments, retirementCode));
                    gene.setSubtype("RETIRED_ANNOTATED_GENE");
                    gene.setAnalysisRun(run);
                }
            }
            sess.commit();
        } catch (Throwable t) {
            sess.rollback();
            if (t instanceof RuntimeException)
                throw (RuntimeException) t;
            else
                throw new ErrorException(t);
        } finally {
            sess.close();
        }
    }

    public void reviewAnnotatedTranscript(String featureId, String reviewer,
            String comment) throws SQLException {
        DbSession sess = new DbSession(connection);
        try {
            AnnotatedTranscript transcript = (AnnotatedTranscript) sess
                    .getById(AnnotatedTranscript.class, featureId);
            transcript.setReviewer(reviewer);
            transcript.setReviewDate(new Date());
            transcript.setComment(comment);
            sess.commit();
        } catch (Throwable t) {
            sess.rollback();
            if (t instanceof RuntimeException)
                throw (RuntimeException) t;
            else
                throw new ErrorException(t);
        } finally {
            sess.close();
        }
    }

    public void insertUtrElement(String sequenceId, int start, int stop,
            String strand, String transcriptId, int transcriptStart,
            int transcriptStop, String subtype, String name)
            throws SQLException {

        DbSession sess = new DbSession(connection);
        try {
            BioSequence seq = (BioSequence) sess.getById(BioSequence.class,
                    sequenceId, false);
            Assert.a(seq != null, "Sequence not found: " + sequenceId);

            AnalysisRunService ars = new AnalysisRunService(sess);
            AnalysisRun run = ars.fetchOrCreateAnalysisRun(seq.getGroup()
                    .getId(), "manual");

            AnnotatedTranscript trans = (AnnotatedTranscript) sess.getById(
                    AnnotatedTranscript.class, transcriptId);

            Number cnt = (Number) sess
                    .get(
                            "select count(*) from UtrElement u where u.sequence.id=? and u.start=? and u.stop=? and u.subtype=?",
                            new Object[] { seq.getId(), new Integer(start),
                                    new Integer(stop), subtype }, true);

            if (cnt.intValue() == 0) {
                UtrElement utr = new UtrElement();
                utr.setSequence(seq);
                utr.setStart(start);
                utr.setStop(stop);
                utr.assignStrand(strand);
                utr.setTranscript(trans);
                utr.setSubtype(subtype);
                utr.setName(name);
                utr.setOntologyTerm("polyA_signal_sequence");
                utr.setAnalysisRun(run);
                sess.save(utr);
            }
            sess.commit();
        } catch (Throwable t) {
            sess.rollback();
            if (t instanceof RuntimeException)
                throw (RuntimeException) t;
            else
                throw new ErrorException(t);
        } finally {
            sess.close();
        }
    }

    public void deleteUtrElement(String featureId) throws SQLException {
        DbSession sess = new DbSession(connection);
        try {
            sess.deleteById(Feature.class, featureId);
            sess.commit();
        } catch (Throwable t) {
            sess.rollback();
            if (t instanceof RuntimeException)
                throw (RuntimeException) t;
            else
                throw new ErrorException(t);
        } finally {
            sess.close();
        }
    }

    private Exon createExon(BioSequence seq, String exonDesc, String strand,
            AnalysisRun run) {
        String[] exonProps = exonDesc.split(",");
        Exon exon = new Exon();
        exon.setSequence(seq);
        exon.setSubtype("ANNOTATED_EXON");
        exon.setStart(Integer.parseInt(exonProps[0]));
        exon.setStop(Integer.parseInt(exonProps[1]));
        exon.assignStrand(strand);
        exon.setType(exonProps[2]);
        exon.setVersion(new Integer(0));
        exon.setAnalysisRun(run);
        return exon;
    }

    private RetiredFeature createRetiredFeature(DbSession sess, Feature f,
            String annotator, String comment, int retirementCode) {
        RetiredFeature result = new RetiredFeature();
        result.setId(f.getId());
        result.setOriginalSubtype(f.getSubtype());
        result.setRetireDate(new Date());
        result.setRetiredBy(annotator);
        result.setComment(comment);
        result.setRetirementCode((RetirementCode) sess.getById(
                RetirementCode.class, "" + retirementCode));

        return result;
    }

    private String join(String[] s) {
        StringBuffer result = new StringBuffer();
        for (int i = 0; (s != null) && (i < s.length); i++) {
            if (i > 0)
                result.append("#");
            result.append(s[i]);
        }
        return result.toString();
    }

    private boolean nullString(String s) {
        return (s == null) || (s.length() == 0);
    }

    private String newAccessionNumber(DbSession sess, Feature f) {
        int accType;
        if (f instanceof Gene)
            accType = GenomeAssembly.GENE_ACCESSION;
        else if (f instanceof Transcript)
            accType = GenomeAssembly.TRANSCRIPT_ACCESSION;
        else
            throw new IllegalArgumentException("Unexpected Feature class: "
                    + f.getClass().getName());

        GenomeAssembly assembly = (GenomeAssembly) sess.getById(
                GenomeAssembly.class, f.getSequenceGroup().getId(), false);
        String accs[] = assembly.getAccessionNumbers(sess, 1, accType);
        return accs[0];
    }

    /**
     * Given the id of a normal transcript, it will convert it into an
     * AnnotatedTranscript in the database.
     * 
     * @param transcriptId
     * @param transcriptTypeId
     * @param createdBy
     * @param modifiedBy
     */
    public void promoteTranscript(String transcriptId, int transcriptTypeId,
            String createdBy, String modifiedBy) {
        DbSession sess = new DbSession(connection);
        sess.setFlushMode(FlushMode.COMMIT);

        // get the parent gene's accession. (Creating a new one if not already
        // set)
        Transcript t = (Transcript) sess
                .getById(Transcript.class, transcriptId);
        Gene gene = t.getGene();
        if (gene.getPrimaryAcc() == null) {
            gene.setPrimaryAcc(newAccessionNumber(sess, gene));
        }
        String geneAccession = gene.getPrimaryAcc();
        try {
            sess.flush();
            sess.clear();
        } catch (Throwable e) {
            fixGene(gene, sess);
            sess.flush();
            sess.clear();
        }
        // just populate all the required not-null fields
        try {
            PreparedStatement stmt = connection
                    .prepareStatement("INSERT INTO ap_annotated_transcript (ap_feature_id, ap_transcript_type_id, ap_gene_accession_number) values (?, ?, ?)");
            stmt.setString(1, transcriptId);
            stmt.setInt(2, transcriptTypeId);
            stmt.setString(3, geneAccession);
            stmt.execute();
            stmt.close();
            stmt = connection
                    .prepareStatement("UPDATE AP_FEATURE set ap_subclass = 'ANTRS' where AP_ID = ?");
            stmt.setString(1, transcriptId);
            stmt.execute();
            stmt.close();
            // looks like I don't need a commit here
            // connection.commit();
        } catch (SQLException ex) {
            throw new ErrorException(
                    "Could not create AnnotatedTranscript record in database",
                    ex);
        }

        // now we should be able to load this as an annotated transcript and do
        // the
        // remaining via hibernate.
        AnnotatedTranscript transcript = (AnnotatedTranscript) sess.getById(
                AnnotatedTranscript.class, transcriptId);
        transcript.setCreator(createdBy);
        transcript.setCreationDate(new Date());
        transcript.setModifier(modifiedBy);
        transcript.setModificationDate(new Date());
        transcript.setSubtype("ANNOTATED_TRANSCRIPT");
        transcript.setType((TranscriptType) sess.getById(TranscriptType.class,
                new Integer(transcriptTypeId), false));
        sess.close(true);
    }

    private Transcript createTranscript(DbSession sess, int mode,
            String sequenceId, int start, int stop, String strand,
            int startCodon, int stopCodon, int transcriptTypeId,
            String transcriptName, String primaryEvidenceFeatureId,
            String transcriptComments, String createdBy, String modifiedBy,
            String[] curationFlagIds, String[] exonList,
            String[] supplementalEvidenceIds, String symbol, String geneId,
            String transcriptId) {
        boolean autoFlag = false;
        System.out.println("REINOS: createTranscript");
        if (curationFlagIds == null)
            curationFlagIds = EMPTY_STRINGS;
        if (exonList == null)
            exonList = EMPTY_STRINGS;
        if (supplementalEvidenceIds == null)
            supplementalEvidenceIds = EMPTY_STRINGS;

        Gene gene = null;
        AnnotatedTranscript transcript = null;
        BioSequence seq = (BioSequence) sess.getById(BioSequence.class,
                sequenceId, false);
        Assert.a(seq != null, "No sequence found: " + sequenceId);
        AnalysisRunService ars = new AnalysisRunService(sess);
        AnalysisRun run = ars.fetchOrCreateAnalysisRun(seq.getGroup().getId(),
                "manual");
        if (nullString(geneId)) {
            gene = new Gene();
            gene.setSequence(seq);
            gene.assignStrand(strand);
            gene.setSubtype("ANNOTATED_GENE");
            gene.setPrimaryAcc(newAccessionNumber(sess, gene));
            gene.setVersion(new Integer(0));
            gene.setAnalysisRun(run);
        } else {
            gene = (Gene) sess.getById(Gene.class, geneId, false);
            Assert.a(gene != null, "No gene found: " + geneId);
            Assert.a(gene.getSequence().equals(seq),
                    "Gene associated with different sequence: "
                            + gene.getSequence().getId());
        }
        gene.setSymbol(symbol);

        if (mode == INSERT_MODE) {
            Assert.a(nullString(transcriptId),
                    "Transcript id must not be specified in insert mode.");
            transcript = new AnnotatedTranscript();
            transcript.setCreator(createdBy);
            transcript.setCreationDate(new Date());
            transcript.setSequence(seq);
            transcript.setAnalysisRun(run);
            transcript.setTranscriptAccessionNum(newAccessionNumber(sess,
                    transcript));
        } else {
            transcript = (AnnotatedTranscript) sess.getById(
                    AnnotatedTranscript.class, transcriptId, false);
            Assert
                    .a(transcript != null, "No transcript found: "
                            + transcriptId);
            Assert.a(seq.equals(transcript.getSequence()),
                    "Transcript associated with different sequence: "
                            + transcript.getSequence().getId());
            transcript.getExons().clear();
            transcript.getSupplementalEvidence().clear();
            transcript.getCurationFlags().clear();
            transcript.setEvidenceFeature(null);
            transcript.setModifier(modifiedBy);
            transcript.setModificationDate(new Date());
            transcript.getGene().removeTranscript(transcript);
        }
        transcript.assignStrand(strand);
        transcript.setStart(start);
        transcript.setStop(stop);
        transcript.setSubtype("ANNOTATED_TRANSCRIPT");

        transcript.setType((TranscriptType) sess.getById(TranscriptType.class,
                new Integer(transcriptTypeId), false));
        transcript.setName(transcriptName);
        if (!nullString(primaryEvidenceFeatureId))
            transcript.setEvidenceFeature((Feature) sess.getById(Feature.class,
                    primaryEvidenceFeatureId));
        transcript.setComments(transcriptComments);
        for (int i = 0; i < curationFlagIds.length; i++) {
            if ("AUTO".equals(curationFlagIds[i])) {
                autoFlag = true;
            } else {
                transcript.addCurationFlag((CurationFlag) sess.getById(
                        CurationFlag.class, curationFlagIds[i]));
            }
        }
        for (int i = 0; i < exonList.length; i++) {
            String exonDesc = exonList[i];
            transcript.addExon(createExon(seq, exonDesc, strand, run));
        }
        for (int i = 0; i < supplementalEvidenceIds.length; i++) {
            transcript.getSupplementalEvidence().add(
                    (Feature) sess.getById(Feature.class,
                            supplementalEvidenceIds[i]));
        }

        transcript.setStartCodon(startCodon);
        transcript.setStopCodon(stopCodon);

        transcript.setGeneAccessionNum(gene.getPrimaryAcc());
        transcript.setVersion(new Integer(0));

        gene.addTranscript(transcript);
        if (autoFlag) {
            fixGene(gene, sess);
        }
        gene.shrinkToTranscripts();
        gene.check();

        return transcript;
    }

    private void fixGene(Gene gene, DbSession sess) {
        Iterator it = gene.getTranscripts().iterator();
        while (it.hasNext()) {
            Transcript t = (Transcript) it.next();
            t.populateCurationFlags(sess);

            // perhaps this should be done after check() ?
            // only decend into objects which are initialized because
            // we don't want to touch more than we have to.
            populateCurationFlagsIfInitialized(t.getEvidenceFeature(), sess);

            Collection c = t.getSupplementalEvidence();
            if (Hibernate.isInitialized(c)) {
                Iterator evidenceIt = c.iterator();
                while (evidenceIt.hasNext()) {
                    populateCurationFlagsIfInitialized((Feature) evidenceIt
                            .next(), sess);
                }
            }
        }
    }

    private void populateCurationFlagsIfInitialized(Feature f, DbSession sess) {
        if (f == null)
            return;

        if (Hibernate.isInitialized(f)) {
            f = (Feature) sess.objectCast(f, f.getNarrowedClass());
            if (f instanceof Transcript) {
                ((Transcript) f).populateCurationFlags(sess);
            }
        }
    }

    private void saveTranscript(DbSession sess, Transcript transcript) {
        System.out.println("REINOS: saveTranscript");
        Gene gene = transcript.getGene();
        try {
            sess.saveOrUpdate(gene);
            sess.saveOrUpdate(transcript);
            sess.commit();
        } catch (Throwable t) {
            log.error(t);
            sess.rollback();
            if (t instanceof RuntimeException)
                throw (RuntimeException) t;
            else
                throw new ErrorException(t);
        } finally {
            sess.close();
        }
    }
}
