// --------------------------------------------------------------------
// 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 tinylib.Point3D;

/**
  Container for ATOM or HETATM entry of pdb file.
**/
class PdbAtom {
  /**
    read a string and returns the atom data.
    The return value is an anonymous container. Components are:

    - ci: chain id
    - het: is hetero atom (bool)
    - rn: residue name
    - ri: residue index
    - at: atom data of `PdbAtom` type
  **/
  static public function readFromText( text:String ):Dynamic {
    var fieldname:String = text.substr( 0, 6 );
    var ishetero:Bool = ( fieldname == "HETATM" );
    var atomindex:Int = Std.parseInt( text.substr( 6, 5 ) );
    // never be empty; if atomname is empty, that atom will be ignored.
    var atomname:String = StringTools.trim( text.substr( 12, 4 ) );
    var resname:String = StringTools.trim( text.substr( 17, 3 ) );
    var resid:Int = Std.parseInt( text.substr( 22, 5 ) );
    var chainid:String = text.substr( 21, 1 );
    var x:Float = Std.parseFloat( text.substr( 30, 8 ) );
    var y:Float = Std.parseFloat( text.substr( 38, 8 ) );
    var z:Float = Std.parseFloat( text.substr( 46, 8 ) );
    var atom:PdbAtom = new PdbAtom( atomname, atomindex, new Point3D( x, y, z ) );
    if ( text.length >= 66 ) {
      var oc:Float = Std.parseFloat( text.substr( 54, 6 ) );
      var bf:Float = Std.parseFloat( text.substr( 60, 6 ) );
      atom.occupancy = oc;
      atom.bfactor = bf;
    }
    // replace element name
    if ( text.length >= 78 ) {
      atom.element = StringTools.trim( text.substr( 76, 2 ) );
    }
    var anon = { ci:chainid, het:ishetero, rn:resname, ri:resid, at:atom.clone() };
    return( anon );
  }

  /**
    guess element name from atom name `n`.
    return value is the name of the element.
  **/
  static public function guessElementName( n:String ):String {
    var len:Int = n.length;
    var str:String = "";
    for ( i in 0 ... len ) {
      var s:String = n.charAt(i);
      // extract only alphabets
      if ( s.toLowerCase() != s.toUpperCase() ) str += s;
    }
    return( PdbAtom.elementName( str ) );
  }

  /**
    returns the distance between two `PdbAtom`s, `a0` and `a1`.
  **/
  static public function getDistance( a0:PdbAtom,
                                      a1:PdbAtom ):Float {
    return( Point3D.getSub( a0.pos, a1.pos ).norm() );
  }

  /**
    Tentatively defined atomic radius; there is no physical/chemical meaning.
    `e` is the element name such as H, C, O, and the return value is the
    radius of the element in angstrom.
  **/
  static public function getRadii( e:String ):Float {
    switch ( e.toUpperCase() ) {
      case "H":
        return( 1.5 );
      case "C", "O", "N", "F":
        return( 2.0 );
      case "NA", "MG":
        return( 2.5 );
      case "P", "S", "CL":
        return( 2.8 );
      case "K", "CAL", "FE", "ZN":
        return( 3.0 );
    }
    return( 2.0 );
  }

  /**
    returns WMXML style xml data of bond between two PdbAtoms, `a0` and `a1`,
    as a string.
  **/
  static public function genBondXml( a0:PdbAtom,
                                     a1:PdbAtom ):String {
    return( "    <BOND pos0=\"" + a0.pos.x + " " + a0.pos.y + " " + a0.pos.z +
                   "\" pos1=\"" + a1.pos.x + " " + a1.pos.y + " " + a1.pos.z +
                   "\" col0=\"" + PdbAtom.getColor( a0.element ) +
                   "\" col1=\"" + PdbAtom.getColor( a1.element ) +
                   "\" />\n" );
  }

  /**
    returns the predefined color of the element.
    `e` is the element name, and the return value is the color name in string
    such as "white", "lime".
  **/
  static public function getColor( e:String ):String {
    switch ( e.toUpperCase() ) {
      case "H":
        return( "white" );
      case "C":
        return( "lime" );
      case "O":
        return( "red" );
      case "N":
        return( "blue" );
      case "F":
        return( "pink" );
      case "NA", "MG":
        return( "magenta" );
      case "P":
        return( "pink" );
      case "S":
        return( "yellow" );
      case "CL":
        return( "green" );
      case "K", "CAL", "FE", "ZN":
        return( "silver" );
    }
    return( "lime" );
  }

  /**
    return element name of atom name `n`.
    This function assumes the name `n` does not contain non-alphabet characters.
    If your input name contains number such as HG2, use `guessElementName`.
  **/
  static public function elementName( n:String ):String {
    while( n.length >= 1 ) {
      switch ( n.toUpperCase() ) {
        case "H": // hydrogen
          return( "H" );
        case "C", "CA": // carbon; CA is assumed to be alpha carbon
          return( "C" );
        case "O": // oxygen
          return( "O" );
        case "N": // nitrogen
          return( "N" );
        case "F": // fluroine
          return( "F" );
        case "NA", "SOD": // sodium (cation)
          return( "NA" );
        case "MG", "MAG": // magnesium (cation)
          return( "MG" );
        case "P": // phosphorus
          return( "P" );
        case "S": // sulfur
          return( "S" );
        case "CL", "CLA": // chloride (anion)
          return( "CL" );
        case "K": // potassium
          return( "K" );
        case "CAL": // calcium
          return( "CA" );
        case "FE": // iron
          return( "FE" );
        case "ZN": // zinc
          return( "ZN" );
        default:
          n = n.substr( 0, n.length - 1 );
      }
    }
    return( "" );
  }

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

  /**
    atom index; 2nd column in ATOM or HETATM record
  **/
  @: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 { index = i; return( index ); }

  /**
    atom name; 3rd column in ATOM or HETATM record
  **/
  @:isVar public var name( get, set ):String;
    /**
      getter of `name`
    **/
    public function get_name():String { return( name ); }
    /**
      setter of `name`
    **/
    public function set_name( n:String ):String { name = n; return( name ); }

  /**
    position of this atom
  **/
  @:isVar public var pos( get, set ):Point3D;
    /**
      getter of `get_pos`
    **/
    public function get_pos():Point3D { return( pos ); }
    /**
      setter of `get_pos`
    **/
    public function set_pos( p:Point3D ):Point3D { pos = p.clone(); return( pos ); }
  /**
    alternate positions; not often used
  **/
  @:isVar public var altpos( get, null ):Array< Point3D >;
    /**
      getter of `altpos`
    **/
    public function get_altpos():Array< Point3D > { return( altpos ); }

  /**
    occupancy in pdb record
  **/
  @:isVar public var occupancy( get, set ):Float;
    /**
      getter of `occupancy`
    **/
    public function get_occupancy():Float { return( occupancy ); }
    /**
      setter of `occupancy`
    **/
    public function set_occupancy( o:Float ):Float { occupancy = o; return( occupancy ); }

  /**
    bfactor in pdb record
  **/
  @:isVar public var bfactor( get, set ):Float;
    /**
      getter of `bfactor`
    **/
    public function get_bfactor():Float { return( bfactor ); }
    /**
      setter of `bfactor`
    **/
    public function set_bfactor( f:Float ):Float { return( bfactor ); }

  /**
    element name of this atom
  **/
  @:isVar public var element( get, set ):String;
    /**
      getter of `element`
    **/
    public function get_element():String { return( element ); }
    /**
      setter of `element`
    **/
    public function set_element( e:String ):String { element = e; return( element ); }

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

  /**
    Constructor.

    - n: atom name
    - i: atom index
    - p: atom position
    - oc: occupancy
    - bf: B factor
    - e: element name
  **/
  public function new( ?n:String = "",
                       ?i:Int = 1,
                       ?p:Point3D,
                       ?oc:Float = 1.0,
                       ?bf:Float = 0.0,
                       ?e:String = "" ) {
    name = n;
    index = i;
    if ( p != null ) pos = p;
    occupancy = oc;
    bfactor = bf;
    element = e;
    if ( name != "" && element == "" ) {
      element = PdbAtom.guessElementName( name );
    }
  }

  /**
    returns clone of `this` instance.
  **/
  public function clone():PdbAtom {
    var ret:PdbAtom = new PdbAtom();
    ret.name = name;
    ret.index = index;
    ret.pos = pos;
    ret.occupancy = occupancy;
    ret.bfactor = bfactor;
    ret.element = element;
    return( ret );
  }

  /**
    add an alternate position to this atom
  **/
  public function addAlternatePos( p:Point3D ):Void {
    altpos.push( p.clone() );
  }

  /**
    returns ATOM record for WMXML xml of this atom.
  **/
  public function genAtomXml():String {
    return( "    <ATOM pos=\"" + pos.x + " " + pos.y + " " + pos.z +
                   "\" color=\"" + PdbAtom.getColor( element ) +
                   "\" />\n" );
  }
}
