/*
   Project: UL

   Copyright (C) 2005 Michael Johnston & Jordi Villa-Freixa

   Author: Michael Johnston

   Created: 2005-05-24 15:03:58 +0200 by michael johnston

   This application 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 application is distributed ithe 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
   Library General Public License for more details.

   You should have received a copy of the GNU General Public
   License along with this library; if not, write to the Free
   Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111 USA.
*/

#include "PDBConfigurationBuilder.h"

@implementation PDBConfigurationBuilder

- (void) _loadPlugin: (NSString*) name
{	
	//search for the requested plugin in the standard location
	//($HOME)/adun/Plugins/Configuration/
	
	NSString* pluginDir;
	NSBundle *pluginBundle;
	NSString *pluginPath;
	NSString *temp;
	Class pluginClass;

	pluginDir = [[ioManager valueForKey: @"applicationDir"]
				stringByAppendingPathComponent: @"Plugins/Configurations"];

	NSDebugLLog(@"PDBConfigurationBuilder", @"Plugin dir is %@. Plugin Name is %@", pluginDir, name);

	//add check to see if bundle actually exists

	pluginBundle = [NSBundle bundleWithPath: 
				[pluginDir stringByAppendingPathComponent: 
				name]];
	if(pluginBundle == nil)
		[NSException raise: NSInvalidArgumentException format: @"Specified plugin does not exist"];	

	NSDebugLLog(@"PDBConfigurationBuilder", @"Plugin Bundle is %@", pluginBundle);
	NSDebugLLog(@"PDBConfigurationBuilder", 
			@"Dynamicaly Loading Plugin from Directory: %@.\n\n", [pluginBundle bundlePath]);

	if(pluginClass = [pluginBundle principalClass])
	{ 
		NSDebugLLog(@"PDBConfigurationBuilder", 
				@"Found plugin Script (plugin=%@).\n", [pluginClass description]);
		plugin = [pluginClass new];

		if(![plugin conformsToProtocol:@protocol(ULPDBConfigurationPlugin)])
			[NSException raise: NSInternalInconsistencyException 
				format: 
				@"Specified plugins (%@) principal class does not conform to ULPDBConfigurationBuilder", 
				[pluginClass description]];	
	}
	else
		[NSException raise: NSInternalInconsistencyException 
			format: @"Specified plugin has no principal class"];

	NSDebugLLog(@"PDBConfigurationBuilder", @"Loaded plugin\n");
}

- (id) _structureObjectForPDB: (NSString*) path
{
	id outputPath, molStruct;
	FILE* file_p;

	outputPath = [[NSUserDefaults standardUserDefaults] stringForKey: @"BuildOutput"]; 
	buildOutput = fopen([outputPath cString], "a");

	if(![[NSFileManager defaultManager] fileExistsAtPath: path])
	{	
		GSPrintf(buildOutput, @"File does not exist!\n");
		[NSException raise: NSInvalidArgumentException
			format: [NSString stringWithFormat: 
			@"Specified file (%@) does not exist", path]];
	}
	
	GSPrintf(buildOutput, @"The pdb file is %@\n\n", path);

	NS_DURING
	{	
		molStruct = [StructureFactory newStructureFromPDBFile: path];
	}
	NS_HANDLER
	{
		[NSException raise: @"ULBuildException" 
			format: @"Unable to create structure object"];
	}
	NS_ENDHANDLER

	fclose(buildOutput);

	return molStruct;
}

/*
Takes an array of amino acid atom names and the residue they
are from. This method replaces any old style names in the
array with the correct version. It returns YES if any replacements
were made, NO otherwise.
note: The old new mapping only will affect certain hydrogens so
it would  be more efficent to just do this on them. 
However we could expand this to deal with other pdb naming errors
involving other naming systems e.g. using IUPAC. This
would be complicated however due to nonunique lables (see below)
*/

- (BOOL) _correctAminoAcidAtomNames: (NSMutableArray*) atomNames
	forResidue: (NSString*) residue
{
	BOOL oldNames = NO;
	int i, numberOfAtoms;
	NSString* currentAtom, *newName;
	NSArray* newNames;
	NSDictionary* nameMapForResidue;
	NSMutableDictionary* replacementMap;

	/*
	 We need to avoid cases where be the same ID exists in both
	 old and new forms but refers to different atoms.
	 Therefore if any atom is old we assume they all are.
	*/
	
	//nameMap maps oldPDBName->newPDBName
	nameMapForResidue = [nameMap objectForKey: residue];

	//if we cant id the residue just return - 
	//we only deal with the 20 standard A.As
	if(nameMap == nil)
		return;

	newNames = [nameMapForResidue allValues];
	numberOfAtoms = [atomNames count];
	for(i=0; i<numberOfAtoms; i++)
	{
		currentAtom = [atomNames objectAtIndex: i];
		if(![newNames containsObject: currentAtom])
		{
			oldNames =YES;
			break;
		}
	}
	
	if(oldNames)
		for(i=0; i<numberOfAtoms; i++)
		{
			currentAtom = [atomNames objectAtIndex: i];
			newName = [nameMapForResidue objectForKey: currentAtom];
			if(newName != nil)
				[atomNames replaceObjectAtIndex: i
					withObject: newName];
		}
}

/*
This method takes a list of amino acids atoms & residues and a mapping
of atoms to residues. It checks if the names are correct and returns
an array of the atom names with the corrections.
*/

- (id) _verifyAminoAcidAtomNames: (NSMutableArray*) atomNames
	forResidues: (NSMutableArray*) residues
	atomsPerResidue: (NSMutableArray*) atomsPerResidue
{
	int startAtom, endAtom, i;
	NSRange residueRange;
	NSString* currentResidue;
	NSMutableArray* residueAtoms, *newAtomNames;

	[buildString appendString: @"\t\tChecking hydrogen naming and correcting if neccessary\n"];
	
	newAtomNames = [NSMutableArray array];
	startAtom = endAtom = 0;
	for(i=0; i<[residues count]; i++)
	{
		currentResidue = [residues objectAtIndex: i];
		endAtom =  [[atomsPerResidue objectAtIndex: i] intValue];
		residueRange = NSMakeRange(startAtom, endAtom);
		residueAtoms = [[atomNames subarrayWithRange: residueRange]
					mutableCopy];
	
		[self _correctAminoAcidAtomNames: residueAtoms
			forResidue: currentResidue];
		
		[newAtomNames addObjectsFromArray: residueAtoms];
		startAtom += endAtom;
	}

	return newAtomNames;
}

- (id) _processChain: (id) chain 
	section: (NSString*) section 
	selection: (NSArray*) selection
	atomMatrix: (ULMatrix*) matrix
{
	int numberOfAtoms;
	id residue, residueName, residueDescription, atom;
	NSEnumerator* residueEnum, *atomEnum;
	NSMutableArray *atomsPerResidue, *sequence, *atomNames, *coordinateArray;
	NSMutableDictionary* result;

	atomNames = [NSMutableArray array];
	atomsPerResidue = [NSMutableArray array];
	sequence = [NSMutableArray array];
	result = [NSMutableDictionary dictionary];
	coordinateArray = [NSMutableArray array];
	
	if([section isEqual: @"Residues"])
		residueEnum = [chain allResidues];
	else if([section isEqual: @"Heterogens"])
		residueEnum = [chain allHeterogens];
	else if([section isEqual: @"Solvent"])
		residueEnum = [chain allSolvent];
	else
		[NSException raise: NSInvalidArgumentException
			format: [NSString stringWithFormat:
		 	@"Invalid chain section %@", section]];

	NSDebugLLog(@"PDBConfigurationBuilder",
			 @"\tProcess chain section %@ with selection %@", section, selection);

	if(selection == nil)
	{
		NSDebugLLog(@"PDBConfigurationBuilder", 
				@"\tSelection is nil. Adding all available residues");
		[buildString appendString: @"\t\tAdding all available residues\n"];

		while(residue = [residueEnum nextObject])
		{
			[atomsPerResidue addObject: [NSNumber numberWithInt: 
				[[[residue allAtoms] allObjects] count]]];
			[sequence addObject:
				[[residue name] stringByTrimmingCharactersInSet: 
					[NSCharacterSet whitespaceCharacterSet]]];
			atomEnum = [residue allAtoms];
			while(atom = [atomEnum nextObject])
			{
				[atomNames addObject: [atom name]];
				[coordinateArray addObject: 
					[NSNumber numberWithDouble: [atom x]]];
				[coordinateArray addObject: 
					[NSNumber numberWithDouble: [atom y]]];
				[coordinateArray addObject: 
					[NSNumber numberWithDouble: [atom z]]];
				[matrix extendMatrixWithRow: coordinateArray];
				[coordinateArray removeAllObjects];
			}
		}
	}	
	else
	{
		[buildString appendString: @"\t\tSearching for selected residues\n"];
		
		while(residue = [residueEnum nextObject])
		{
			residueName = [residue name];
			residueDescription = [residue description];
			NSDebugLLog(@"PDBConfigurationBuilder",
				@"\tChecking residue %@", residueDescription);
			if([selection containsObject: residueDescription])
			{	
				NSDebugLLog(@"PDBConfigurationBuilder",
					@"\t\tAdding %@", residueDescription);
				[buildString appendFormat: @"\t\tAdding %@\n", residueDescription];

				[atomsPerResidue addObject: [NSNumber numberWithInt: 
					[[[residue allAtoms] allObjects] count]]];
				[sequence addObject:
					[[residue name] stringByTrimmingCharactersInSet: 
						[NSCharacterSet whitespaceCharacterSet]]];
				atomEnum = [residue allAtoms];
				while(atom = [atomEnum nextObject])
				{
					[atomNames addObject: [atom name]];
					[coordinateArray addObject: 
						[NSNumber numberWithDouble: [atom x]]];
					[coordinateArray addObject: 
						[NSNumber numberWithDouble: [atom y]]];
					[coordinateArray addObject: 
						[NSNumber numberWithDouble: [atom z]]];
					[matrix extendMatrixWithRow: coordinateArray];
					[coordinateArray removeAllObjects];
				}
			 }
		}
	}	
	
	NSDebugLLog(@"PDBConfigurationBuilder",
			@"\tSection processed:\nSequence %@\nAtoms %@\nAtoms per Residue %@",
			sequence, atomNames, atomsPerResidue);
	[buildString appendFormat: @"\t%d residues\n\t%d atoms\n",
			[sequence count], [atomNames count]];

	
	//It we are dealing with amino acids (residues)
	//replace atomNames with the checked array returned
	//by _verifyAtomNames:forResidues:atomsPerResidue
	if([section isEqual: @"Residues"])
		atomNames = [self _verifyAminoAcidAtomNames: atomNames 
			forResidues: sequence
			atomsPerResidue: atomsPerResidue];

	[result setObject: sequence forKey: @"Sequence"];
	[result setObject: atomNames forKey: @"AtomNames"];
	[result setObject: atomsPerResidue forKey: @"AtomsPerResidue"];

	return result;
}

- (id) _configurationObjectFromPDBStructure: (NSDictionary*) options
{
	int i;
	NSEnumerator *chainEnum;
	NSMutableArray *atomNames, *sequences, *sequence, *atomsPerResidue;
	NSDictionary* result;
	NSString* name;
	id chain, selectedSections, selectedChainOptions;
	NSMutableArray* selectedChains;
	id configuration, matrix;

	selectedChains = [options valueForKey: @"Selection"];

	if([selectedChains count] == 0)
		[NSException raise: @"ULBuildException"
			format: @"No chains were selected"];

	[buildString appendFormat: @"Selected Chains %@\n\n", selectedChains];
	chainEnum = [structure allChains];
	atomNames = [NSMutableArray array];
	atomsPerResidue = [NSMutableArray array];
	sequences = [NSMutableArray array];
	matrix = [ULMatrix new];

	i = 1;
	while(chain = [chainEnum nextObject])
	{
		//the chain names in the options selections are in the
		//form "Chain $CHAINCODE" - 

		name = [chain name];
		name = [name stringByTrimmingCharactersInSet: 
				[NSCharacterSet whitespaceCharacterSet]];
		if(name == nil || [name isEqual: @""])
			name = [NSString stringWithFormat: @"Chain %d", i];
		else
			name = [NSString stringWithFormat: @"Chain %@", name];

		NSDebugLLog(@"PDBConfigurationBuilder", @"Chain Name = %@", name);

		//check if this chain was selected

		if([selectedChains containsObject: name])
		{	
			[buildString appendFormat: @"Processing %@\n", name];

			selectedChainOptions = [options valueForKey: name];
			selectedSections = [selectedChainOptions
						valueForKey: @"Selection"];
			sequence = [NSMutableArray array];	
			GSPrintf(buildOutput, @"Chain %@ containing %d residues\n", 
						[chain name], [chain countResidues]);
			GSPrintf(buildOutput, @"The sequence is: %@\n\n", [chain get3DSequence]);	
			
			//check residues

			if([selectedSections containsObject: @"Residues"])
			{
				[buildString appendString: @"\tResidues Selected\n"];
				result = [self _processChain: chain 
						section: @"Residues" 
						selection: nil
						atomMatrix: matrix];
				[sequence addObjectsFromArray:
					[result objectForKey: @"Sequence"]];		
				[atomsPerResidue addObjectsFromArray:
					[result objectForKey: @"AtomsPerResidue"]];
				[atomNames addObjectsFromArray:
					[result objectForKey: @"AtomNames"]];
			}

			//check heterogens
			
			if([selectedSections containsObject: @"Heterogens"])
			{	
				[buildString appendString: @"\tHeterogens Selected\n"];
				result = [self _processChain: chain 
						section: @"Heterogens" 
						selection: [selectedChainOptions 
							 valueForKeyPath: @"Heterogens.Selection"]
						atomMatrix: matrix];
				[sequence addObjectsFromArray:
					[result objectForKey: @"Sequence"]];		
				[atomsPerResidue addObjectsFromArray:
					[result objectForKey: @"AtomsPerResidue"]];
				[atomNames addObjectsFromArray:
					[result objectForKey: @"AtomNames"]];
			}	

			//check solvent
			
			if([selectedSections containsObject: @"Solvent"])
			{
				[buildString appendString: @"\tSolvent Selected\n"];
				result = [self _processChain: chain 
						section: @"Solvent" 
						selection: [selectedChainOptions
							 valueForKeyPath: @"Solvent.Selection"]
						atomMatrix: matrix];
				[sequence addObjectsFromArray:
					[result objectForKey: @"Sequence"]];		
				[atomsPerResidue addObjectsFromArray:
					[result objectForKey: @"AtomsPerResidue"]];
				[atomNames addObjectsFromArray:
					[result objectForKey: @"AtomNames"]];
			}	

			if([sequence count] != 0)
				[sequences addObject: sequence];
			
			[buildString appendFormat: @"Completed %@\n\n", name];
		}
		else
			NSDebugLLog(@"PDBConfigurationBuilder", @"%@ was not selected", name);

		i++;
	}

	if([sequences count] == 0)
		[NSException raise: @"ULBuildException"
			format: @"No residues were selected in any chain"];

	NSDebugLLog(@"PDBConfigurationBuilder", @"\nProcess pdb - Building configuration");

	configuration = [NSMutableDictionary dictionaryWithCapacity:1];
	[configuration setObject: matrix forKey: @"Coordinates"];
	[configuration setObject: atomNames forKey: @"AtomNames"];
	[configuration setObject: sequences forKey: @"Sequences"];
	[configuration setObject: atomsPerResidue forKey: @"AtomsPerResidue"];
	if([structure pdbcode] != nil)
		[configuration setObject: [structure pdbcode] forKey: @"SystemName"];
	else
		[configuration setObject: @"None" forKey: @"SystemName"];

	return configuration;
}

/***************

Public Methods

****************/

- (id) init
{
	return [self initWithMoleculeAtPath: nil];
}

- (id) initWithMoleculeAtPath: (NSString*) path
{
	NSString* pathExtension;
	NSString* nameMapFile;

	NSDebugLLog(@"PDBConfigurationBuilder", 
			@"Initialising configuration builder - path is %@", path);

	if(self = [super initWithMoleculeAtPath: path])
	{
		ioManager = [ULIOManager appIOManager];
		if(path != nil)
		{
			pathExtension = [[moleculePath pathExtension] lowercaseString];
			if([[moleculePath pathExtension] isEqual: @"pdb"])
				structure = [self _structureObjectForPDB: path];
			else
				return nil;

			if(structure == nil)
				return nil;

			moleculePath = [path retain];
			[structure retain];
		} 
		else
			moleculePath = nil;

		pluginName =  nil;

		//load pdb name map
		
		nameMapFile = [[[NSBundle mainBundle] resourcePath]
				 stringByAppendingPathComponent: @"pdbNameMap.plist"];	
		nameMap = [NSDictionary dictionaryWithContentsOfFile: nameMapFile];
		[nameMap retain];
	}		
		
	return self;
}

- (NSMutableDictionary*) buildOptions
{
	int i;
	id chain, name, residue;
	id mainOptions;
	NSString* optionsFile;
	NSEnumerator* chainEnum, *enumerator;
	NSMutableArray* chainSelection, *nameArray ;
	NSMutableDictionary* chainOptions;

	if(structure == nil)
		return nil;

	optionsFile = [[[NSBundle mainBundle]
			 resourcePath]
			 stringByAppendingPathComponent: @"pdbBuilderOptions.plist"];	

	//The options file contains the structure for one chain

	mainOptions = [NSMutableDictionary dictionaryWithContentsOfFile: optionsFile];
	[mainOptions removeObjectForKey: @"ChainA"];

	chainEnum = [structure allChains];
	chainSelection = [NSMutableArray array];
	i = 1;
	while((chain = [chainEnum nextObject]))
	{
		chainOptions = [NSMutableDictionary dictionaryWithContentsOfFile: optionsFile];
		chainOptions = [chainOptions valueForKey: @"ChainA"];
	
		NSDebugLLog(@"PDBConfigurationBuilder", @"Chain options %@", chainOptions);
	
		if([chain countResidues] == 0)
		{	
			[chainOptions removeObjectForKey: @"Residues"];	
			[[chainOptions objectForKey: @"Selection"] removeAllObjects];
		}		

		if([chain countSolvent] != 0)
		{
			nameArray = [NSMutableArray array];
			enumerator = [chain allSolvent];
			while(residue = [enumerator nextObject])
				[nameArray addObject: [residue description]];

			[chainOptions setValue: nameArray
					forKeyPath: @"Solvent.Choices"];
		}
		else
			[chainOptions removeObjectForKey: @"Solvent"];			
		
		if([chain countHeterogens] != 0)
		{
			nameArray = [NSMutableArray array];
			enumerator = [chain allHeterogens];
			while(residue = [enumerator nextObject])
				[nameArray addObject: [residue description]];

			[chainOptions setValue: nameArray 
				forKeyPath: @"Heterogens.Choices"];
		}
		else
			[chainOptions removeObjectForKey: @"Heterogens"];			

		name = [chain name];
		name = [name stringByTrimmingCharactersInSet: 
				[NSCharacterSet whitespaceCharacterSet]];
		if(name == nil || [name isEqual: @""])
			name = [NSString stringWithFormat: @"Chain %d", i];
		else
			name = [NSString stringWithFormat: @"Chain %@", name];

		//check that there was a least something in this chain
		//since sometimes libmoltalk invents phantom chains
		if([chainOptions count] > 2)
		{
			[chainSelection addObject: name];
			[mainOptions setObject: chainOptions 
				forKey: name];
		}

		i++;
	}

	//The options display code cant handle muliple default selections.
	//So if we add all the chains here only the first one will appear to be selected
	//in the initial display but they all will be selected when building
	//To avoid this we only select the first one for the moment

	[[mainOptions objectForKey: @"Selection"] addObject: 
		[chainSelection objectAtIndex: 0]];

	return mainOptions;
}

- (id) buildConfiguration: (NSDictionary*) options 
		error: (NSError**) buildError
		userInfo: (NSString**) buildInfo
{
	id configuration, outputPath;
	
	[buildString release];
	buildString = [[NSMutableString stringWithCapacity: 1] retain];
	*buildInfo = buildString; 	

	NSDebugLLog(@"PDBConfigurationBuilder", @"Options are %@\n", options);
	NSDebugLLog(@"PDBConfigurationBuilder",  @"Structure object %@", structure);
	
	outputPath = [[NSUserDefaults standardUserDefaults] stringForKey: @"BuildOutput"]; 
	buildOutput = fopen([outputPath cString], "a");

	[buildString appendFormat: @"\nConfiguration File: %@\n", moleculePath];
	configuration = [self _configurationObjectFromPDBStructure: options];

	fclose(buildOutput);
	
	NSDebugLLog(@"PDBConfigurationBuilder", @"%@", configuration);
	NSDebugLLog(@"PDBConfigurationBuilder", @"Printing matrix");
	NSDebugLLog(@"PDBConfigurationBuilder", @"No rows %d", 
		[[configuration objectForKey: @"Coordinates"] Rows]);

	NSDebugLLog(@"PDBConfigurationBuilder", @"Complete");
	[buildString appendString: @"Completed configuration build\n"];

	return configuration;	
}

- (NSMutableDictionary*) optionsForPlugin
{
	return [plugin options];
}

- (void) applyPlugin: (NSDictionary*) options
{
	if(plugin != nil)
	{
		NSDebugLLog(@"PDBConfigurationBuilder", @"Calling main on plugin %@", plugin);
		[plugin manipulatePDBStructure: structure 
			userOptions: options];
	}
}

- (void) loadPlugin: (NSString*) name
{
	NSDebugLLog(@"PDBConfigurationBuilder", @"Plugin name is %@", name);

	if([pluginName isEqual: @"None"])
		plugin = nil;
	else
	{
		[self _loadPlugin: name];
		[pluginName release];
		pluginName = name;
		[pluginName retain];
	}
}

- (NSString*) currentPlugin
{
	return pluginName;
}

- (void) setCurrentMolecule: (NSString*) path
{
	NSDebugLLog(@"PDBConfigurationBuilder", @"Path %@", path);
	
	if(path ==  nil)
		[self removeCurrentMolecule];
	else if([[path pathExtension] isEqual: @"pdb"])
	{
		//\note change this to structureForCurrentMolecule
		NSDebugLLog(@"PDBConfigurationBuilder", @"Creating stucture for %@", path);
		[structure release];
		structure = [self _structureObjectForPDB: path];
		[structure retain];
		NSDebugLLog(@"PDBConfigurationBuilder", @"Structure created successfully");
		[moleculePath release];
		moleculePath = [path retain]; 
	}
	else
		[NSException raise: NSInvalidArgumentException
			format: [NSString stringWithFormat:
			 @"File at %@ is not have a valid extension"]];

}

- (void) removeCurrentMolecule
{
	[moleculePath release];
	[structure release];
	moleculePath = nil;
	structure = nil;
}

- (NSString*) currentMoleculePath
{
	return moleculePath;
}

@end
