// K-3D
// Copyright (c) 1995-2004, Timothy M. Shead
//
// Contact: tshead@k-3d.com
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation; either
// version 2 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.
//
// You should have received a copy of the GNU General Public
// License along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

/** \file
		\brief Implements the PDBReader K-3D object, which reads Protein Database files.  The standard is defined at http://www.rcsb.org/pdb/docs/format/pdbguide2.2/guide2.2_frame.html
		\author Ed Millard (emillard@direcway.com)
*/

#include <k3dsdk/classes.h>
#include <k3dsdk/color.h>
#include <k3dsdk/file_helpers.h>
#include <k3dsdk/high_res_timer.h>
#include <k3dsdk/ideletable.h>
#include <k3dsdk/idocument.h>
#include <k3dsdk/ifile_format.h>
#include <k3dsdk/igeometry_read_format.h>
#include <k3dsdk/imaterial.h>
#include <k3dsdk/imaterial_collection.h>
#include <k3dsdk/imouse_event_observer.h>
#include <k3dsdk/iobject.h>
#include <k3dsdk/iobject_collection.h>
#include <k3dsdk/material.h>
#include <k3dsdk/module.h>
#include <k3dsdk/plugins.h>
#include <k3dsdk/property.h>
#include <k3dsdk/string_modifiers.h>
#include <k3dsdk/transform.h>
#include <k3dsdk/utility.h>
#include <k3dsdk/vectors.h>

#include <boost/filesystem/fstream.hpp>
#include <boost/filesystem/path.hpp>

#include <iostream>
#include <map>

namespace
{

/////////////////////////////////////////////////////////////////////////////
// pdb_reader_implementation

class pdb_reader_implementation :
	public k3d::ifile_format,
	public k3d::igeometry_read_format,
	public k3d::ideletable
{
public:
	unsigned long priority()
	{
		return 0;
	}

	bool query_can_handle(const boost::filesystem::path& FilePath)
	{
		return "pdb" == k3d::file_extension(FilePath);
	}

	bool pre_read(k3d::idocument&, const boost::filesystem::path& FilePath)
	{
		return true;
	}

	bool read_options(k3d::idocument& Document, const boost::filesystem::path& FilePath)
	{
		return true;
	}

	bool read_file(k3d::idocument& Document, const boost::filesystem::path& FilePath);

	k3d::iplugin_factory& factory()
	{
		return get_factory();
	}

	static k3d::iplugin_factory& get_factory()
	{
		static k3d::plugin_factory<k3d::application_plugin<pdb_reader_implementation>, k3d::interface_list<k3d::igeometry_read_format> > factory(
			k3d::uuid(0x43488899, 0x635d46d4, 0xa5aef51d, 0x0c4d019a),
			"PDBReader",
			"Protein Database ( .pdb )",
			"");

		return factory;
	}
};

class material_factory
{
public:
	material_factory() :
		m_carbon(0),
		m_nitrogen(0),
		m_oxygen(0),
		m_hydrogen(0),
		m_sulphur(0),
		m_iron(0)
	{
	}

	k3d::imaterial* carbon(k3d::idocument& Document)
	{
		if(!m_carbon)
			m_carbon = create_material(Document, "Carbon", k3d::color(0, 1, 0));
		return m_carbon;
	}

	k3d::imaterial* nitrogen(k3d::idocument& Document)
	{
		if(!m_nitrogen)
			m_nitrogen = create_material(Document, "Nitrogen", k3d::color(0, 0, 1));
		return m_nitrogen;
	}

	k3d::imaterial* oxygen(k3d::idocument& Document)
	{
		if(!m_oxygen)
			m_oxygen = create_material(Document, "Oxygen", k3d::color(1, 0, 0));
		return m_oxygen;
	}

	k3d::imaterial* hydrogen(k3d::idocument& Document)
	{
		if(!m_hydrogen)
			m_hydrogen = create_material(Document, "Hydrogen", k3d::color(1, 1, 1));
		return m_hydrogen;
	}

	k3d::imaterial* sulphur(k3d::idocument& Document)
	{
		if(!m_sulphur)
			m_sulphur = create_material(Document, "Sulphur", k3d::color(1, 1, 0));
		return m_sulphur;
	}

	k3d::imaterial* iron(k3d::idocument& Document)
	{
		if(!m_iron)
			m_iron = create_material(Document, "Iron", k3d::color(1, 0, 1));
		return m_iron;
	}


private:
	k3d::imaterial* create_material(k3d::idocument& Document, const std::string& Name, const k3d::color Color)
	{
		static k3d::idocument_plugin_factory* const material_factory = dynamic_cast<k3d::idocument_plugin_factory*>(k3d::plugin(k3d::classes::RenderManMaterial()));
		return_val_if_fail(material_factory, 0);

		k3d::iobject* const material_object = material_factory->create_plugin(Document);
		return_val_if_fail(material_object, 0);

		material_object->set_name(Name);
		assert_warning(k3d::set_property_value(*material_object, "color", Color));

		Document.objects().add_objects(k3d::make_collection<k3d::iobject_collection::objects_t>(material_object));

		return dynamic_cast<k3d::imaterial*>(material_object);
	}

	k3d::imaterial* m_carbon;
	k3d::imaterial* m_nitrogen;
	k3d::imaterial* m_oxygen;
	k3d::imaterial* m_hydrogen;
	k3d::imaterial* m_sulphur;
	k3d::imaterial* m_iron;
};

const double kRadiusCarbon		= 1.20; // Angstroms
const double kRadiusNitrogen	= 1.05;
const double kRadiusOxygen		= 1.00;
const double kRadiusHydrogen	= 0.70;
const double kRadiusSulphur		= 1.25;
const double kRadiusIron		= 1.72;

bool pdb_reader_implementation::read_file(k3d::idocument& Document, const boost::filesystem::path& FilePath)
{
	// Try to open the input file ...
	boost::filesystem::ifstream file(FilePath);
	if(!file.good()) {
		std::cerr << __PRETTY_FUNCTION__ << ": error opening [" << FilePath.native_file_string() << "]" << std::endl;
		return_val_if_fail(0, false);
	}

	// This vector holds a record for each object, its parent and depth in
	// the hierarchy.  All the objects are added to the document and hierarchy
	// at the end of the file load to improve file loading performance
	k3d::iobject_collection::objects_t records;

	// Atom factory
	k3d::idocument_plugin_factory* const atom_factory =
		dynamic_cast<k3d::idocument_plugin_factory*>(k3d::plugin(k3d::uuid(0x45588899, 0x635d46d4, 0xa5aef51d, 0x0c4d019a)));

	// Parent molecule
	k3d::idocument_plugin_factory* const molecule_factory =
		dynamic_cast<k3d::idocument_plugin_factory*>(k3d::plugin(k3d::uuid(0x45588855, 0x635d46d4, 0xa5aef51d, 0x0c4d019a)));
	k3d::iobject* molecule = molecule_factory->create_plugin(Document);
	records.insert(molecule);

	// Materials
	material_factory materials;

	// Chain
	k3d::iobject* chain = NULL;
	std::string	currentChainID(" ");

	// Residue
	k3d::iobject* residue = NULL;
	std::string	currentResidueString;


	while(!file.eof()) {

		// Grab one line at a time ...
		std::string linebuffer;
		k3d::getline(file, linebuffer);

		// Skip empty lines ...
		if(0 == linebuffer.size())
			continue;

		// Extract a record type ...
		std::istringstream stream(linebuffer);

		size_t recordLength = 6;
		if(recordLength > linebuffer.size()) {
			recordLength = linebuffer.size();
		}
		std::string recordtype;
		stream >> recordtype;

		// Skip empty lines ...
		if(0 == recordtype.size())
			continue;

		// Process material file(s) ...
		if(recordtype == "ATOM"			// Coordinate
		|| recordtype == "HETATM") {	// Heterogen

			std::istringstream string(linebuffer.substr(6, 5));
			int serialNumber;
			string >> serialNumber;

			std::string atomName	= linebuffer.substr(12, 4);
			std::string residueName	= linebuffer.substr(17, 3);
			std::string chainID		= linebuffer.substr(21, 1);

			std::string residueSequence(linebuffer.substr(22, 4));

			std::istringstream positionString(linebuffer.substr(30, 24));
			k3d::vector3 position;
			positionString >> position;

			std::istringstream occupancyString(linebuffer.substr(54, 6));
			double occupancy;
			occupancyString >> occupancy;

			std::istringstream tempString(linebuffer.substr(60, 6));
			double tempFactor;
			tempString >> tempFactor;

#if 0
			std::string segmentID(linebuffer.substr(72, 4));
			std::string element;
			if(linebuffer.size() >= 78) {
				element = linebuffer.substr(76, 2);
			}
#endif

			//std::cerr << "atom " << serialNumber << " " << atomName << " res " << residueName;
			//std::cerr << " chain " << chainID << " resSeq " << residueSequence;
			//std::cerr << " " << position << std::endl;

			if(chainID != currentChainID) {
				chain = molecule_factory->create_plugin(Document);
				chain->set_name(chainID);
				records.insert(chain);
				currentChainID = chainID;
			}

			residueName += residueSequence;
			if(residueName != currentResidueString) {
				residue = molecule_factory->create_plugin(Document);
				residue->set_name(residueName);

				if(chain) {
					records.insert(residue);
				}else{
					records.insert(residue);
				}

				currentResidueString = residueName;
			}
			k3d::iobject* atom = atom_factory->create_plugin(Document);
			atom->set_name(atomName);
			int depth;
			if(chain) depth = 3;
			else depth = 2;

			records.insert(atom);

			k3d::set_position(*atom, position);
			double radius = 1.0;
			k3d::imaterial_collection* material_collection = dynamic_cast<k3d::imaterial_collection*>(atom);
			const char* name = atomName.c_str();
			if(name[0] == ' ') {
				name++;
			}
			if(name[0] == 'C') {
				radius = kRadiusCarbon;
				material_collection->set_material(materials.carbon(Document));
			}else if(name[0] == 'H'){
				radius = kRadiusHydrogen;
				material_collection->set_material(materials.hydrogen(Document));
			}else if(name[0] == 'O'){
				radius = kRadiusOxygen;
				material_collection->set_material(materials.oxygen(Document));
			}else if(name[0] == 'N'){
				radius = kRadiusNitrogen;
				material_collection->set_material(materials.nitrogen(Document));
			}else if(name[0] == 'S'){
				radius = kRadiusSulphur;
				material_collection->set_material(materials.sulphur(Document));
			}else if(name[0] == 'F' && name[1] == 'E'){
				radius = kRadiusIron;
				material_collection->set_material(materials.iron(Document));
			}else{
				printf("unimplemented %s\n",name);
			}
			k3d::set_property_value(*atom, "radius", radius);

		}else if(recordtype == "CONECT") {	// Connectivity
			const unsigned int kFieldWidth = 5;
			const unsigned int kMaxAtoms = 11;

			int serialNumbers[kMaxAtoms];
			unsigned int col = 6;
			for(unsigned int i = 0 ; i < kMaxAtoms ; ++i) {
				if(linebuffer.size() >= (col + kFieldWidth)) {
					std::istringstream string(linebuffer.substr(col, kFieldWidth));
					if(string >> serialNumbers[i]) {
					}else{
						serialNumbers[i] = 0;
					}
				}else{
					serialNumbers[i] = 0;
				}
				col += kFieldWidth;
			}
		}else if(recordtype == "REMARK") {	// Header
		}else if(recordtype == "HEADER") {	// Header
		}else if(recordtype == "OBSLTE") {	// Header
		}else if(recordtype == "CAVEAT") {	// Header
		}else if(recordtype == "TITLE" || recordtype == "COMPND") {	// Header
			if(molecule == NULL) {
				std::string name;
				stream >> name;
				molecule->set_name(name);
				printf("molecule %s\n", name.c_str());
			}
		}else if(recordtype == "SOURCE") {	// Header
		}else if(recordtype == "KEYWDS") {	// Header
		}else if(recordtype == "EXPDTA") {	// Header
		}else if(recordtype == "AUTHOR") {	// Header
		}else if(recordtype == "REVDAT") {	// Header
		}else if(recordtype == "SPRSDE") {	// Header
		}else if(recordtype == "JRNL") {	// Header

		}else if(recordtype == "MODEL") {	// Coordinate
		}else if(recordtype == "SIGATM") {	// Coordinate
		}else if(recordtype == "ANISOU") {	// Coordinate
		}else if(recordtype == "SIGUIJ") {	// Coordinate
		}else if(recordtype == "TER") {		// Coordinate
			// Terminate chain
			chain = NULL;
			currentChainID = " ";
		}else if(recordtype == "ENDMDL") {	// Coordinate

		}else if(recordtype == "ORIGXn") {	// Transform
		}else if(recordtype == "SCALEn") {	// Transform
		}else if(recordtype == "MTRIXn") {	// Transform
		}else if(recordtype == "TVECT") {	// Transform

		}else if(recordtype == "HETNAM") {	// Heterogen
		}else if(recordtype == "HETSYN") {	// Heterogen
		}else if(recordtype == "FORMUL") {	// Heterogen

		}else if(recordtype == "DBREF") {	// Primary structure
		}else if(recordtype == "SEQADV") {	// Primary structure
		}else if(recordtype == "SEQRES") {	// Primary structure
		}else if(recordtype == "MODRES") {	// Primary structure

		}else if(recordtype == "HELIX") {	// Secondary structure
		}else if(recordtype == "SHEET") {	// Secondary structure
		}else if(recordtype == "TURN") {	// Secondary structure

		}else if(recordtype == "CRYST1") {	// Crystallography

		}else if(recordtype == "SITE") {	// Misc

		}else if(recordtype == "SSBOND") {	// Connect note
		}else if(recordtype == "LINK") {	// Connect note
		}else if(recordtype == "HYDBND") {	// Connect note
		}else if(recordtype == "SLTBRG") {	// Connect note
		}else if(recordtype == "CISPEP") {	// Connect note

		}else if(recordtype == "MASTER") {	// Bookkeeping
		}else if(recordtype == "END") {		// Bookkeeping
		}
	}

	Document.objects().add_objects(records);

	return true;
}

} // namespace

namespace libk3dchem
{

k3d::iplugin_factory& pdb_reader_factory()
{
	return pdb_reader_implementation::get_factory();
}

} // namespace libk3dchem