/*
 * Copyright (c) Medical Research Council 2002. All rights reserved.
 *
 * Permission to use, copy, modify and distribute this software and its
 * documentation for any purpose is hereby granted without fee, provided that
 * this copyright and notice appears in all copies.
 *
 * This file was written as part of the Staden Package at the MRC Laboratory
 * of Molecular Biology, Hills Road, Cambridge, CB2 2QH, United Kingdom.
 *
 * MRC disclaims all warranties with regard to this software.
 *
 */


#include <cassert>
#include <cstring>                    // For strcpy(), memset()
#include <cstdio>                     // For std::printf(), std::sprintf()
#include <new>                        // For std::bad_alloc
#include <algorithm>                  // For std::max()
#include <list.hpp>                   // For list template
#include <align.hpp>                  // For Alignment object
#include <trace.hpp>                  // For Trace object
#include <muttag.hpp>                 // For MutTag object
#include <peakcall.hpp>               // For PeakCall object
#include <tagarray.hpp>               // For TagArray object
#include <tracealign.hpp>             // For helpers
#include <tracealign_cache.hpp>       // For TraceAlignCache object
#include <tracealign_preprocess.hpp>  // For TraceAlignPreprocessor object



// #define VERBOSE_DEBUG 



/**
   Initialises an empty tracealign_t structure.
*/
void TraceAlignInit( tracealign_t* ta )
{
   assert(ta != NULL);
   std::memset( ta, 0, sizeof(tracealign_t) );
   ta->ResultString    = new char[ 512 ];
   ta->ResultString[0] = 0;
   ta->Initialised     = 1;
}



/**
   Frees up all memory allocated by the trace alignment algorithm
   in the tracealign_t structure.
*/
void TraceAlignDestroy( tracealign_t* ta )
{
   assert(ta != NULL);
   assert(ta->Initialised);
   try
   {
      // Delete all data
      TraceAlignDestroyCache( ta );
      TraceAlignDestroyResults( ta );
      delete [] ta->ResultString;
   }
   catch(...)
   {
      // Shouldn't happen, but we musn't throw exceptions outside dll boundary
      assert(0);
   }
}



/**
   Sets the specified reference trace. The clip points are in base units.
   The caller retains ownership of the Read structure.
*/
void TraceAlignSetReference( tracealign_t* ta, mutlib_strand_t d, Read* r, int ql, int qr )
{
   assert(r != NULL);
   assert(ta != NULL);
   assert(ta->Initialised);
   ta->Reference[d].ClipL  = ql;
   ta->Reference[d].ClipR  = qr;
   ta->Reference[d].Trace  = r;
   ta->Reference[d].Strand = d;
   ta->Reference[d].New    = 1;
}



/**
   Sets the input traces. The clip points are in base units. The caller retains
   ownership of the Read structure.
*/
void TraceAlignSetInput( tracealign_t* ta, mutlib_strand_t d, Read* r, int ql, int qr )
{
   assert(r != NULL);
   assert(ta != NULL);
   assert(ta->Initialised);
   ta->Input.ClipL  = ql;
   ta->Input.ClipR  = qr;
   ta->Input.Trace  = r;
   ta->Input.Strand = d;
   ta->Input.New    = 1;
}



/**
   Returns the result code generated by the Execute() function. This can be
   used to query the result later without having to save the value returned
   by Execute().
*/
mutlib_result_t TraceAlignGetResultCode( tracealign_t* ta )
{
   assert(ta != NULL);
   assert(ta->Initialised);
   return ta->ResultCode;
}



/**
   Returns a read-only user friendly error message corresponding to the
   last result code. It may contain more useful information such as a
   filename or the particular intput that was invalid. This can be displayed
   to the user if required.
*/
const char* TraceAlignGetResultString( tracealign_t* ta )
{
   assert(ta != NULL);
   assert(ta->Initialised);
   return ta->ResultString;
}



/**
   Returns the aligned traces as Read structures. Ownership of the Read
   structures are retained by the trace alignment algorithm. If 'l' or 'r'
   is not null, the left and right clip points in base numbers are returned
   as well.
*/
Read* TraceAlignGetAlignment( tracealign_t* ta, mutlib_input_t i, int* l, int* r )
{
   assert(ta != NULL);
   assert(ta->Initialised);
   if( l != 0 )
      *l = ta->Alignment[i].ClipL;
   if( r != 0 )
      *r = ta->Alignment[i].ClipR;
   return ta->Alignment[i].Trace;
}



/**
   Executes the trace alignment algorithm.
*/
mutlib_result_t TraceAlignExecute( tracealign_t* ta )
{
    enum { STATE_INITIALISE, STATE_VALIDATE_INPUT, STATE_UPDATE_CACHE,
           STATE_PREPROCESS_INPUT, STATE_SEQUENCE_ALIGN, STATE_SEQUENCE_OVERLAP,
           STATE_SEQUENCE_MAP, STATE_QUANTISE_ENVELOPE, STATE_ESTIMATE_BANDSIZE,
           STATE_ENVELOPE_ALIGN, STATE_TRACE_ALIGN, STATE_EXIT };
    mutlib_strand_t         Strand = MUTLIB_STRAND_FORWARD; // silence warning
    Alignment               Aligner;
    int                     AlignedSeqOverlap[2];
    DNAArray<char>          RefSeq;
    TraceAlignPreprocessor* RefData = NULL;
    Trace                   RefTrace;
    SimpleArray<char>       RefEnvelope;
    SimpleArray<char>       RefEnvelopeAligned;
    int                     RefOverlapBases[2];
    int                     RefOverlapSamples[2];
    DNAArray<char>          InputSeq;
    TraceAlignPreprocessor  InputData;
    Trace                   InputTrace;
    SimpleArray<char>       InputEnvelope;
    SimpleArray<char>       InputEnvelopeAligned;
    int                     InputOverlapBases[2];
    int                     InputOverlapSamples[2];
    const int               L           = 0;
    const int               R           = 1;
    const int               QLOWER      = 1;
    const int               QLEVELS     = 126;
    const char              QPAD_SYMBOL = QLEVELS + QLOWER;
    const int               MIN_OVERLAP = 10;
    bool                    Done        = false;
    int                     State       = STATE_INITIALISE;
    TraceAlignCache*        Cache       = 0;
    int                     BandSize    = 0;
    assert(ta != NULL);
    assert(ta->Initialised);
    try
    {
        while(!Done)
        {
            switch(State)
            {
                case STATE_INITIALISE:
                    // Destroy old results
                    TraceAlignDestroyResults( ta );
                    Strand = ta->Input.Strand;
                    // Initialise trace objects
                    RefTrace.Wrap( ta->Reference[Strand].Trace, false );
                    InputTrace.Wrap( ta->Input.Trace, false );
                    State = STATE_VALIDATE_INPUT;
                    break;



                case STATE_VALIDATE_INPUT:
                    // Check input values
                    if( TraceAlignValidateInput(ta) )
                        State = STATE_EXIT;
                    else
                        State = STATE_UPDATE_CACHE;
                    break;



                case STATE_UPDATE_CACHE:
                    // Initialise cache if not done
                    if( !ta->Cache )
                        ta->Cache = static_cast<void*>( new TraceAlignCache );
                    Cache = static_cast<TraceAlignCache*>(ta->Cache);
                    if( Cache->AlignmentMatrix.Rows() <= 0 )
                        Cache->CreateAlignmentMatrix( QLEVELS+QLOWER+1, QLEVELS, QLOWER );
                    // If new wildtype has been supplied
                    if( ta->Reference[Strand].New )
                    {
                        // Preprocess wildtype and cache it
                        Cache->RefData[Strand].Flush();
                        Cache->RefData[Strand].PreprocessTrace( RefTrace, true );
                        ta->Reference[Strand].New = 0;

                    }
                    RefData = &Cache->RefData[Strand];
                    State    = STATE_PREPROCESS_INPUT;
                    break;



                case STATE_PREPROCESS_INPUT:
                    // Preprocess input trace
                    InputData.PreprocessTrace( InputTrace );
                    State = STATE_SEQUENCE_ALIGN;
                    break;



                case STATE_SEQUENCE_ALIGN: {
                    // Align sequences first so we can determine which parts of the
                    // trace to align - trace alignment takes a long time using DP!
                    char* wseq = &(RefTrace.Raw()->base[ ta->Reference[Strand].ClipL ]);
                    char* iseq = &(InputTrace.Raw()->base[ ta->Input.ClipL ]);
                    int   wlen = ta->Reference[Strand].ClipR - ta->Reference[Strand].ClipL - 1;
                    int   ilen = ta->Input.ClipR - ta->Input.ClipL - 1;
                    Aligner.InputSequence( 0, wseq, wlen );
                    Aligner.InputSequence( 1, iseq, ilen );
                    Aligner.EdgeScore( Alignment::EDGE_SCORE_RIGHT );
                    Aligner.Execute( Alignment::ALGORITHM_NORMAL );
                    #ifdef VERBOSE_DEBUG
                    Aligner.DumpToFile( "sequence.align.txt" );
                    #endif
                    RefSeq.Create( Aligner.OutputSequence(0), Aligner.OutputSequenceLength(0) );
                    InputSeq.Create( Aligner.OutputSequence(1), Aligner.OutputSequenceLength(1) );
                    State = STATE_SEQUENCE_OVERLAP;
                    break; }



                case STATE_SEQUENCE_OVERLAP:
                    // Checks that there is sufficient sequence overlap for an alignment
                    // operation to make any logical sense!
                    AlignedSeqOverlap[L] = Aligner.OutputSequenceLeftOverlap(0);
                    AlignedSeqOverlap[R] = Aligner.OutputSequenceRightOverlap(0);
                    if( AlignedSeqOverlap[R] - AlignedSeqOverlap[L] <= MIN_OVERLAP )
                    {
                        std::sprintf( ta->ResultString, "Insufficient sequence overlap to compute "
                                     "trace alignment for %s.\n", ta->Input.Trace->trace_name );
                        ta->ResultCode = MUTLIB_RESULT_INSUFFICIENT_OVERLAP;
                        State = STATE_EXIT;
                        break;

                    }
                    State = STATE_SEQUENCE_MAP;
                    break;



                case STATE_SEQUENCE_MAP:
                    // Maps the sequence alignment start/end points back onto the original
                    // trace. Two mappings are made, one for bases, the other for samples.
                    RefOverlapBases[L]   = RefSeq.GetOriginalPosition( AlignedSeqOverlap[L] );
                    RefOverlapBases[R]   = RefSeq.GetOriginalPosition( AlignedSeqOverlap[R] );
                    InputOverlapBases[L]  = InputSeq.GetOriginalPosition( AlignedSeqOverlap[L] );
                    InputOverlapBases[R]  = InputSeq.GetOriginalPosition( AlignedSeqOverlap[R] );
                    RefOverlapBases[L]  += ta->Reference[Strand].ClipL;
                    RefOverlapBases[R]  += ta->Reference[Strand].ClipL;
                    InputOverlapBases[L] += ta->Input.ClipL;
                    InputOverlapBases[R] += ta->Input.ClipL;
                    // At this point we have the original sequence zero-based base numbers
                    // for the overlapping regions, now we get the corresponding sample points
                    RefOverlapSamples[L]  = RefTrace.BasePosition( RefOverlapBases[L] );
                    RefOverlapSamples[R]  = RefTrace.BasePosition( RefOverlapBases[R] );
                    InputOverlapSamples[L] = InputTrace.BasePosition( InputOverlapBases[L] );
                    InputOverlapSamples[R] = InputTrace.BasePosition( InputOverlapBases[R] );
                    State = STATE_QUANTISE_ENVELOPE;
                    break;



                case STATE_QUANTISE_ENVELOPE: {
                    // Quantise the trace envelopes of the overlapping regions. Note
                    // that the quantised envelope indices will not be relative to the
                    // orignal trace unless you add xxxOverlapSamples[L] to them.
                    int QUPPER[2];
                    RefData->Envelope().Range( RefOverlapSamples[L], RefOverlapSamples[R] );
                    InputData.Envelope().Range( InputOverlapSamples[L], InputOverlapSamples[R] );
                    #ifdef VERBOSE_DEBUG
                    RefData->Envelope().SaveAs( "qenvelope.txt" );
                    InputData.Envelope().SaveAs( "qenvelope.txt", true );
                    #endif
                    QUPPER[0] = RefData->Envelope().Max();
                    QUPPER[1] = InputData.Envelope().Max();
                    QUPPER[0] = std::max( QUPPER[0], QUPPER[1] );
                    TraceAlignQuantiseEnvelope( RefData->Envelope(), RefEnvelope, QLEVELS, QLOWER, QUPPER[0] );
                    TraceAlignQuantiseEnvelope( InputData.Envelope(), InputEnvelope, QLEVELS, QLOWER, QUPPER[0] );
                    State = STATE_ESTIMATE_BANDSIZE;
                    break; }



                case STATE_ESTIMATE_BANDSIZE: {
                    // Estimate an optimal bandsize for the alignment routine to ensure
                    // it runs as fast as possible. We do this by calculating the maximum
                    // absolute difference between the sample vectors, stepping through
                    // one base at a time in the aligned sequence and keeping track of
                    // the maximum difference in samples. Note that the bandsize specified
                    // here is half the width of actual band used in the DP algorithm.
                    int diff;
                    int w, i;
                    int max_diff = 0;
                    const int MIN_BANDSIZE = 10;
                    for( int n=AlignedSeqOverlap[L]; n<=AlignedSeqOverlap[R]; n++ )
                    {
                       w = RefSeq.GetOriginalPosition(n) + ta->Reference[Strand].ClipL;
                       i = InputSeq.GetOriginalPosition(n) + ta->Input.ClipL;
                       w = RefTrace.BasePosition(w) - RefOverlapSamples[L];
                       i = InputTrace.BasePosition(i) - InputOverlapSamples[L];
                       diff = (w >= i) ? w-i : i-w;
                       if( diff > max_diff )
                           max_diff = diff;
                    }
                    BandSize = max_diff + MIN_BANDSIZE;
                    #ifdef VERBOSE_DEBUG
                    std::printf("Alignment Bandsize = %d+%d\n", max_diff, MIN_BANDSIZE );
                    #endif
                    State = STATE_ENVELOPE_ALIGN;
                    break; }



                case STATE_ENVELOPE_ALIGN: {
                    // Align the quantised envelopes of the overlapping region.
                    Aligner.PadSymbol( QPAD_SYMBOL );
                    Aligner.GapPenalty( 0, QLEVELS );
                    Aligner.EdgeScore( Alignment::EDGE_SCORE_RIGHT );
                    Aligner.Matrix( Cache->AlignmentMatrix.Raw(), Cache->AlignmentMatrix.Rows(), false );
                    Aligner.BandSize( BandSize );
                    Aligner.InputSequence( 0, RefEnvelope.Raw(), RefEnvelope.Length() );
                    Aligner.InputSequence( 1, InputEnvelope.Raw(), InputEnvelope.Length() );
                    Aligner.Execute( Alignment::ALGORITHM_NORMAL );
                    RefEnvelopeAligned.Wrap( Aligner.OutputSequence(0), Aligner.OutputSequenceLength(0) );
                    InputEnvelopeAligned.Wrap( Aligner.OutputSequence(1), Aligner.OutputSequenceLength(1) );
                    #ifdef VERBOSE_DEBUG
                    char FileName[256];
                    std::printf("Alignment Score = %f\n", Aligner.OutputScore() );
                    std::sprintf( FileName, "qenvelope.align.txt" );
                    Aligner.DumpToFile( FileName, true );
                    #endif
                    State = STATE_TRACE_ALIGN;
                    break; }



                case STATE_TRACE_ALIGN: {
                    // Align the original traces using the envelope alignment algorithm
                    // output, then insert the bases.
                    Trace RefInterpolated;
                    Trace InputInterpolated;
                    RefInterpolated.Create( RefEnvelopeAligned.Length(), RefOverlapBases[R]-RefOverlapBases[L]+1, RefTrace.Name() );
                    InputInterpolated.Create( InputEnvelopeAligned.Length(), InputOverlapBases[R]-InputOverlapBases[L]+1, InputTrace.Name() );
                    TraceAlignInterpolate( QPAD_SYMBOL, RefEnvelopeAligned, RefTrace, RefOverlapSamples[L], RefInterpolated );
                    TraceAlignInterpolate( QPAD_SYMBOL, InputEnvelopeAligned, InputTrace, InputOverlapSamples[L], InputInterpolated );
                    TraceAlignInsertBases( QPAD_SYMBOL, RefEnvelopeAligned, RefTrace, RefInterpolated, RefOverlapBases );
                    TraceAlignInsertBases( QPAD_SYMBOL, InputEnvelopeAligned, InputTrace, InputInterpolated, InputOverlapBases );
                    ta->Alignment[MUTLIB_INPUT].Trace           = InputInterpolated.Raw();
                    ta->Alignment[MUTLIB_INPUT].ClipL           = InputOverlapBases[L];
                    ta->Alignment[MUTLIB_INPUT].ClipR           = InputOverlapBases[R];
                    ta->Alignment[MUTLIB_INPUT_REFERENCE].Trace = RefInterpolated.Raw();
                    ta->Alignment[MUTLIB_INPUT_REFERENCE].ClipL = RefOverlapBases[L];
                    ta->Alignment[MUTLIB_INPUT_REFERENCE].ClipR = RefOverlapBases[R];
                    RefInterpolated.AutoDestroy( false );
                    InputInterpolated.AutoDestroy( false );
                    #ifdef VERBOSE_DEBUG
                    RefInterpolated.SaveAs( "alignment1.ztr" );
                    InputInterpolated.SaveAs( "alignment2.ztr" );
                    #endif
                    State = STATE_EXIT;
                    break; }



                case STATE_EXIT:
                    // Quit
                    Done = true;
                    break;
            }
        }
    }
    catch( std::bad_alloc& )
    {
        // Memory allocation error
        ta->ResultCode = MUTLIB_RESULT_OUT_OF_MEMORY;
        std::strcpy( ta->ResultString, "Not enough memory available to complete the operation.\n" );
    }
    catch(...)
    {
        // Unknown exceptions
        ta->ResultCode = MUTLIB_RESULT_UNEXPECTED_EXCEPTION;
        std::strcpy( ta->ResultString, "An unexpected fatal exception has occurred, please "
                     "report the details to staden-package@mrc-lmb.cam.ac.uk.\n" );
    }



    // Exit
    return ta->ResultCode;
}