// --------------------------------------------------------------------
// 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.PdbChain;
import wmutil.CellIndex;

/**
  Container for a MODEL in pdb file.
  If there are no MODELs in pdb file, at least one MODEL exists and is used
  in this program.
**/

class PdbModel {
  /**
    chains which belong to this MODEL
  **/
  @:isVar public var chains( get, null ):Map< String, PdbChain >;
    /**
      getter of `chains`
    **/
    public function get_chains():Map< String, PdbChain > { return( chains ); }

  /**
    index of this MODEL
  **/
  @:isVar public var index( get, set ):Int;
    /**
      getter of `index`
    **/
    public function get_index():Int { return( index ); }
    /**
      setter of `index`
    **/
    public function set_index( i:Int ):Int { return( index ); }

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

  /**
    Constructor. You can set MODEL number by `i`.
  **/
  public function new( ?i:Int = 1 ) {
    chains = new Map< String, PdbChain >();
    index = i;
  }

  /**
    add an PdbAtom `atom` to this MODEL.

    - cid: chain id
    - resname: residue name
    - resid: residue index
    - atom: a PdbAtom
  **/
  public function addAtom( cid:String,
                           resname:String,
                           resid:Int,
                           atom:PdbAtom ) {
    
    if ( !chains.exists( cid ) ) {
      chains.set( cid, new PdbChain( cid ) );
    }
    var mychain:PdbChain = chains.get( cid );
    mychain.addAtom( resname, resid, atom );
  }

  /**
    generate WMXML style xml data of this MODEL and return it as a string.
  **/
  public function genXml():String {
    var ret:String = "  <SCENE>\n";
    for ( chain in chains ) {
      ret += chain.genXml();
    }
    ret += genXmlHeteroes();
    ret += "  </SCENE>\n";
    return( ret );
  }

  /**
    return only NON-CHAIN, hetero atom part of XML data as a string.
  **/
  public function genXmlHeteroes():String {
    return( genXmlNotChains( gatherHeteroes() ) );
  }

  /**
    generate non-CHAIN xml data; ATOMs and BONDs
  **/
  public function genXmlNotChains( atoms:Array< PdbAtom > ):String {
    // bond counter for each atom
    // atoms without bonds are shown in sphere, others are shown in sticks
    // ad hoc offset of the array index
    var offset:Int = pdb.PdbGridParams.offset;
    var indexes = new Array();
    // grid size for cell division
    var grid:Float = pdb.PdbGridParams.grid;
    // elements of a position must be larger than - (offset*grid)
    var minx:Int = 99999999; // ad hoc
    var maxx:Int = 0;
    var miny:Int = 99999999;
    var maxy:Int = 0;
    var minz:Int = 99999999;
    var maxz:Int = 0;
    // cell division
    for ( i in 0 ... atoms.length ) {
      var indx:Int = Std.int( atoms[i].pos.x / grid ) + offset;
      var indy:Int = Std.int( atoms[i].pos.y / grid ) + offset;
      var indz:Int = Std.int( atoms[i].pos.z / grid ) + offset;
      // terrible
      if ( indexes[indx] == null ) {
        indexes[indx] = new Array< Array< Array< Int > > >();
      }
      if ( indexes[indx][indy] == null ) {
        indexes[indx][indy] = new Array< Array< Int > >();
      }
      if ( indexes[indx][indy][indz] == null ) {
        indexes[indx][indy][indz] = new Array< Int >();
      }
      indexes[indx][indy][indz].push( i );
      // nothing to say
      minx = Std.int( Math.min( minx, indx ) );
      miny = Std.int( Math.min( miny, indy ) );
      minz = Std.int( Math.min( minz, indz ) );
      maxx = Std.int( Math.max( maxx, indx ) );
      maxy = Std.int( Math.max( maxy, indy ) );
      maxz = Std.int( Math.max( maxz, indz ) );
    }
    return( generateXmlAtomsBonds( minx, miny, minz, maxx, maxy, maxz,
                                   atoms, indexes ) );
  }

  /**
    generate ATOMs and BONDs for heteroes.
  **/
  public function generateXmlAtomsBonds( minx:Int,
                                         miny:Int,
                                         minz:Int,
                                         maxx:Int,
                                         maxy:Int,
                                         maxz:Int,
                                         atoms:Array< PdbAtom >,
                                         indexes:Array< Array< Array< Array< Int > > > > ):String {
    var ret:String = "";
    // number of bonds for each atom
    var nbonds:Array< Int > = new Array< Int >();
    for ( i in 0 ... atoms.length ) nbonds.push(0);
    // four-dimensional array is so terrible
    var cells:Array< CellIndex > = CellIndex.adjCells();
    for ( i in minx ... maxx + 1 ) {
      if ( indexes[i] == null ) continue;
      for ( j in miny ... maxy + 1 ) {
        if ( indexes[i][j] == null ) continue;
        for ( k in minz ... maxz + 1 ) {
          if ( indexes[i][j][k] == null ) continue;
          var num:Int = indexes[i][j][k].length;
          // intra-cell
          for ( i0 in 0 ... num ) {
            var myid0 = indexes[i][j][k][i0];
            var myatom0 = atoms[myid0];
            var myradii0:Float = PdbAtom.getRadii( myatom0.element );
            for ( i1 in i0 + 1 ... num ) {
              var myid1 = indexes[i][j][k][i1];
              var myatom1 = atoms[myid1];
              var dist = PdbAtom.getDistance( myatom0, myatom1 );
              var myradii1:Float = PdbAtom.getRadii( myatom1.element );
              if ( dist < ( myradii0 + myradii1 ) / 2.0 ) {
                // there is a bond between these two atoms
                ret += PdbAtom.genBondXml( myatom0, myatom1 );
                nbonds[myid0] += 1;
                nbonds[myid1] += 1;
              }
            }
          }
          // inter-cell
          for ( cell in cells ) {
            var t = new CellIndex( i + cell.i, j + cell.j, k + cell.k );
            if ( indexes[t.i] == null ) continue;
            if ( indexes[t.i][t.j] == null ) continue;
            if ( indexes[t.i][t.j][t.k] == null ) continue;
            var num1:Int = indexes[t.i][t.j][t.k].length;
            // this loop will be merged into the for loop above?
            for ( i0 in 0 ... num ) {
              var myid0 = indexes[i][j][k][i0];
              var myatom0 = atoms[myid0];
              var myradii0:Float = PdbAtom.getRadii( myatom0.element );
              for ( i1 in 0 ... num1 ) {
                var myid1 = indexes[t.i][t.j][t.k][i1];
                var myatom1 = atoms[myid1];
                var dist = PdbAtom.getDistance( myatom0, myatom1 );
                if ( dist < ( myradii0 + PdbAtom.getRadii( myatom1.element ) ) / 2 ) {
                  ret += PdbAtom.genBondXml( myatom0, myatom1 );
                  nbonds[myid0] += 1;
                  nbonds[myid1] += 1;
                }
              }
            }
          }
        }
      }
    }
    for ( i in 0 ... atoms.length ) {
      if ( nbonds[i] == 0 ||
           atoms[i].element == "NA" ||
           atoms[i].element == "MG" ||
           atoms[i].element == "ZN" ||
           atoms[i].element == "FE" ) {
        ret += atoms[i].genAtomXml();
      }
    }
    return( ret );
  }

  /**
    gather hetero atoms from all chains
  **/
  public function gatherHeteroes():Array< PdbAtom > {
    var atoms:Array< PdbAtom > = new Array< PdbAtom >();
    for ( chain in chains ) {
      for ( residue in chain.residues ) {
        if ( !residue.isAmino() ) {
          for ( atom in residue.atoms ) {
            atoms.push( atom );
          }
        }
      }
    }
    return( atoms );
  }

  /**
    initialize secondary structure type.
    COIL for amino acids and NONE for others.
  **/
  public function initializeSecondary():Void {
    for ( chain in chains ) {
      for ( residue in chain.residues ) {
        if ( residue.isAmino() ) {
          residue.secondary = SecondaryType.COIL;
        } else {
          residue.secondary = SecondaryType.NONE;
        }
      }
    }
  }

  /**
    returns secondary structure information as a string.
    See the same name function in `Pdb.hx` for details.
  **/
  public function genSecondaryString():String {
    var ret:String = "";
    for ( chain in chains ) {
      ret += chain.genSecondaryString();
    }
    return( ret );
  }
}
