// --------------------------------------------------------------------
// wm3d - A Flash Molecular Viewer
//
// Copyright (c) 2011-2013, tamanegi (tamanegi@users.sourceforge.jp)
// All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
// --------------------------------------------------------------------

package pdb;

import pdb.PdbResidue;
import pdb.BackboneHB;
import pdb.RamaParams;

// hash key is (residue name)+(residue id). (ex. GLN102, ATP9912)

/**
  Container of a chain in pdb file
**/

class PdbChain {
  /**
    number of interpolation points used when creating smooth chain in wm3d
  **/
  static public var interpolation:Int = 3;

  // #####################################################################
  /**
    residues in this chain, where the key of Map is residue index
  **/
  @:isVar public var residues( get, null ):Map< Int, PdbResidue >;
    /**
      getter of `residues`
    **/
    public function get_residues():Map< Int, PdbResidue > { return( residues ); }

  /**
    name of this chain
  **/
  @:isVar public var name( get, set ):String;
    public function get_name():String { return( name ); }
    public function set_name( n:String ):String { return( name ); }

  /**
    minimum residue index in this chain
  **/
  @:isVar public var min_index( get, set ):Int;
    /**
      getter of `min_index`
    **/
    public function get_min_index():Int { return( min_index ); }
    /**
      setter of `min_index`
    **/
    public function set_min_index( i:Int ):Int {
      min_index = i;
      return( min_index );
    }
  /**
    maximum residue index in this chain
  **/
  @:isVar public var max_index( get, set ):Int;
    /**
      getter of `max_index`
    **/
    public function get_max_index():Int { return( max_index ); }
    /**
      setter of `max_index`
    **/
    public function set_max_index( i:Int ):Int {
      max_index = i;
      return( max_index );
    }

  // ####################################################################

  /**
    Constructor. You can set chain name by `n`.
  **/
  public function new( ?n:String = " " ) {
    residues = new Map< Int, PdbResidue >();
    name = n;
    min_index = -1;
    max_index = -1;
  }

  /**
    add an PdbAtom `atom `to the `resid` in this chain.
    `resname` is used only when creating new residue.
  **/
  public function addAtom( resname:String,
                           resid:Int,
                           atom:PdbAtom ):Void {
    // resid is the hash key
    if ( !residues.exists( resid ) ) {
      residues.set( resid, new PdbResidue( resname, resid ) );
    }
    residues.get( resid ).addAtom( atom );
    if ( min_index < 0 ) {
      min_index = resid;
      max_index = resid;
    } else {
      min_index = Std.int( Math.min( min_index, resid ) );
      max_index = Std.int( Math.max( max_index, resid ) );
    }
  }

  /**
    calculate secondary structure information
  **/
  public function calcSecondaryStructures():Void {
    // add H atom if does not exist
    // note: cannot add to N-terminal residue
    for ( i in min_index + 1 ... max_index + 1 ) {
      if ( residues.get(i) == null || residues.get(i-1) == null ) continue;
      residues.get(i).genH( residues.get(i-1) );
    }
    // calculated phi and psi angles
    for ( i in min_index + 1 ... max_index ) {
      var resm = residues.get(i-1);
      var res = residues.get(i);
      var resp = residues.get(i+1);
      if ( resm == null || resp == null ) continue;
      var a_ca = res.getAtom( "CA" );
      var a_c = res.getAtom( "C" );
      var a_n = res.getAtom( "N" );
      if ( a_ca == null || a_c == null || a_n == null ) continue;
      // calc phi(N-CA)
      var a_cm = resm.getAtom( "C" );
      if ( a_cm != null ) {
        res.rama.calcPhi( a_cm.pos, a_n.pos, a_ca.pos, a_c.pos );
      }
      // calc psi(CA-C)
      var a_np = resp.getAtom( "N" );
      if ( a_np != null ) {
        res.rama.calcPsi( a_n.pos, a_ca.pos, a_c.pos, a_np.pos );
      }
    }
    // hydrogen bonds
    var hb = new Map< String, BackboneHB >();
    for ( ires in residues ) {
      for ( jres in residues ) {
        var i:Int = ires.index;
        var j:Int = jres.index;
        var key = getMapKey( i, j );
        if ( hb.exists( key ) ) continue;
        hb.set( key, new BackboneHB( ires, jres ) );
      }
    }
    __searchAlphaHelix( hb );
    __search3_10Helix( hb );
    __searchPiHelix( hb );
    __searchBetaStrand( hb );
    __searchTurn( hb );
  }

  // get map key for hydrogen bond map
  private function getMapKey( i:Int,
                              j:Int ):String {
    return( Std.string( Math.min(i,j) + " " + Math.max(i,j) ) );
  }

  private function __searchAlphaHelix( hb:Map< String, BackboneHB > ) {
    // search i - i+4 hydrogen bond
    for ( i in min_index ... max_index + 1 ) {
      if ( residues.get(i) == null ) continue;
      // check i - i+4 pair
      if ( __checkAlphaHelix( i, hb ) ) {
        // check i+1 - i+5 pair
        if ( __checkAlphaHelix( i+1, hb ) ) {
          // two successive component of helix found; they are alpha helix
          for ( k in i ... i + 6 ) {
            var res = residues.get(k);
            if ( res != null ) {
              res.secondary = pdb.SecondaryType.HELIX_ALPHA;
            }
          }
        }
      }
    }
  }

  private function __checkAlphaHelix( index:Int,
                                      hb:Map< String, BackboneHB > ):Bool {
    if ( residues.get(index) == null || residues.get(index+4) == null ) return( false );
    // first, check hydrogen bond
    var key = getMapKey( index, index+4 );
    if ( !hb.get( key ).bonded( HBParams.threshold_hA ) ) return( false );
    // get phi and psi angles of these four residues
    var value = 0.0;
    for ( j in index ... index + 5 ) {
      var res = residues.get(j);
      if ( res == null ) continue;
      value += RamaParams.getValueHelixAlpha( res.rama );
    }
    if ( value >= RamaParams.threshold_hA ) return( true );
    return( false );
  }

  private function __search3_10Helix( hb:Map< String, BackboneHB > ) {
    // not implemented now
  }

  private function __searchPiHelix( hb:Map< String, BackboneHB > ) {
    // not implemented now
  }

  private function __searchBetaStrand( hb:Map< String, BackboneHB > ) {
    for ( i in min_index ... max_index + 1 ) {
      if ( residues.get(i) == null ) continue;
      for ( j in i + 1 ... max_index + 1 ) {
        if ( residues.get(j) == null ) continue;
        // check i - j pair
        if ( __checkBetaStrand( i, j, hb ) ) {
          var n = 1;
          // check i+1 - j+1 pair; parallel beta
          while( __checkBetaStrand( i+n, j+n, hb, (n%2!=0) ) ) {
            n = n + 1;
          }
          if ( __checkBetaStrand( i-1, j-1, hb, true ) ) {
            for ( m in -1 ... n ) {
              var mres = residues.get(i+m);
              if ( mres != null ) mres.secondary = SecondaryType.STRAND;
              var nres = residues.get(j+m);
              if ( nres != null ) nres.secondary = SecondaryType.STRAND;
            }
          } else if ( n > 1 ) {
            for ( m in 0 ... n ) {
              var mres = residues.get(i+m);
              if ( mres != null ) mres.secondary = SecondaryType.STRAND;
              var nres = residues.get(j+m);
              if ( nres != null ) nres.secondary = SecondaryType.STRAND;
            }
          }
          // check i+1 - j-1 pair; anti-parallel beta
          while( __checkBetaStrand( i+n, j-n, hb, (n%2!=0) ) ) {
            n = n + 1;
          }
          if ( __checkBetaStrand( i-1, j+1, hb, true ) ) {
            for ( m in -1 ... n ) {
              var mres = residues.get(i+m);
              if ( mres != null ) mres.secondary = SecondaryType.STRAND;
              var nres = residues.get(j-m);
              if ( nres != null ) nres.secondary = SecondaryType.STRAND;
            }
          } else if ( n > 1 ) {
            for ( m in 0 ... n ) {
              var mres = residues.get(i+m);
              if ( mres != null ) mres.secondary = SecondaryType.STRAND;
              var nres = residues.get(j-m);
              if ( nres != null ) nres.secondary = SecondaryType.STRAND;
            }
          } else {
            var ires = residues.get(i);
            if ( ires != null ) ires.secondary = SecondaryType.BRIDGE;
            var jres = residues.get(j);
            if ( jres != null ) jres.secondary = SecondaryType.BRIDGE;
          }
        }
      }
    }
  }

  private function __checkBetaStrand( indexi:Int,
                                      indexj:Int,
                                      hb:Map< String, BackboneHB >,
                                      ?skipHB:Bool = false ):Bool {
    // first, check hydrogen bond
    var key = getMapKey( indexi, indexj );
    if ( !skipHB ) {
      if ( !hb.get(key).bonded( HBParams.threshold_s ) ) return( false );
    }
    // get phi and psi angles of these four residues
    var value = 0.0;
    value += RamaParams.getValueStrand( residues.get(indexi).rama );
    value += RamaParams.getValueStrand( residues.get(indexj).rama );
    if ( value >= RamaParams.threshold_s ) {
      return( true );
    }
    return( false );
  }

  private function __searchTurn( hb:Map< String, BackboneHB > ) {
    // not implemented now
  }

  /**
    assign secondary structure `sec` to all residues
  **/
  public function assignSec( sec:PdbSecondary ):Void {
    for ( i in sec.init ... sec.last ) {
      var myres = residues.get(i);
      if ( myres != null ) {
        myres.secondary = sec.sectype;
      }
    }
  }

  /**
    returns WMXML xml as a string
  **/
  public function genXml():String {
    var ret:String = "    <CHAIN N=\"" + Std.string(PdbChain.interpolation) + "\">\n";
    for ( i in min_index ... max_index + 1 ) {
      var myres = residues.get(i);
      if ( myres == null ) continue;
      var str:String = myres.genPoint();
      if ( str == "" ) {
        // residues without alpha carbons are assumed to be hetero
        myres.secondary = SecondaryType.NONE;
      } else {
        ret += str;
      }
    }
    ret += __genXmlSecondary();
    ret += "    </CHAIN>\n";
    return( ret );
  }

  private function __genXmlSecondary():String {
    var ret:String = "";
    var secname:Int = SecondaryType.NONE;
    var initnum = 0;
    var prevnum:Int = -1;
    // TODO: sophisticate logic
    for ( i in min_index ... max_index + 1 ) {
      var myres = residues.get(i);
      // nothing to do for non-existing residue
      if ( myres == null ) continue;
      // initialize; first residue and some other cases
      if ( secname == SecondaryType.NONE ) {
        secname = myres.secondary;
        initnum = myres.index;
      } else {
        // 1. secondary structure type of current residue is different from prev
        // 2. current index is not next to the prev one
        // 3. current residue is terminal residue
        // 4. current residue is the last residue
        if ( secname != myres.secondary ||
             prevnum != myres.index - 1 ||
             i == max_index ) {
          var mysecname:Int = secname;
          if ( i == max_index ) {
            prevnum = myres.index;
            mysecname = myres.secondary;
          }
          if ( mysecname != SecondaryType.NONE ) {
            ret += "      <" + SecondaryType.toString( mysecname ) +
                   " init=\"" + initnum +
                   "\" last=\"" + prevnum + "\" />\n";
          }
          initnum = myres.index;
          // need to be re-initialized
          if ( i == max_index ) {
            secname = SecondaryType.NONE;
          } else {
            secname = myres.secondary;
          }
        }
      }
      prevnum = myres.index;
    }
    return( ret );
  }

  /**
    returns secondary structure information of this chain.
    See `Pdb.hx` for details about notations.
  **/
  public function genSecondaryString():String {
    var ret:String = "";
    var prev:PdbResidue = residues.get(min_index);
    for ( i in min_index ... max_index + 1 ) {
      var res:PdbResidue = residues.get(i);
      if ( res == null && prev != null ) ret += " ";
      if ( res != null ) {
        ret += SecondaryType.toChar( res.secondary );
      }
      prev = res;
    }
    return( ret );
  }
}
