#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on Wed Feb  3 12:24:35 2016

@author: tam10
"""

import numpy as np
import csv
import argparse

def get_atom_pair_distances(input_file,pair_list,output_file='output.csv'):
    
    
    if input_file.split('.')[-1]=='log': 
        try:
            from irc2xyz import irc2xyz
            irc2xyz(input_file)
            xyz_file=''.join(input_file.split('.')[:-1])+'.xyz'
        except ImportError:
            raise ImportError("irc2xyz.py not found. Make sure irc2xyz.py is in the same directory as pairdistances.py")
    elif input_file.split('.')[-1]=='xyz':
        xyz_file=input_file
    else:
        raise IOError('Input must be an xyz file or irc log file: \''+xyz_file+'\'')
        
        
    if output_file.split('.')[1]!='csv': raise IOError('Output must be a csv file: \''+output_file+'\'')
    
    with open(xyz_file,'r') as irc: irc_lines=[l.strip('\n').split() for l in irc.readlines()]
    mol_index,mol_collection=-1,[]
    
    for l in irc_lines:
        if l[0].isdigit():                                 #First line represents number of atoms in molecule
            size,atom_index,symbol_list=int(l[0]),0,[]
            mol_collection.append([])                      #This allows appending to a list within mol_collection
            mol_index+=1
        else:
            if atom_index!=0:
                mol_collection[mol_index].append([float(c) for c in l[1:4]])     #Append coordinates
                symbol_list.append([l[0]])                                       #Append atom symbols
                assert size>=atom_index,'Error parsing file: \''+xyz_file+'\''   #Check there aren't more atoms than there should be
            atom_index+=1
    
    distance_list=[]
    for pair in pair_list:
        for atom_index in pair: assert size>=atom_index,'Atom index \''+str(atom_index)+'\' greater than size of molecule ('+str(size)+'). Check input'
        title_string=[symbol_list[pair[0]-1][0]+str(pair[0])+'-'+symbol_list[pair[1]-1][0]+str(pair[1])]                                #Title string in the format "C1-C2"
        distances=[[round(np.linalg.norm(a-m[pair[1]-1]),3) for i,a in enumerate(m) if i==pair[0]-1] for m in np.array(mol_collection)] #Using numpy to subtract lists and calculate distances
        distance_list.append(title_string+[d[0] for d in distances])

    with open(output_file,'w') as output:
        csv_writer=csv.writer(output)
        for index,line in enumerate(distance_list[0]):
            csv_writer.writerow([index+0]+[[l for i,l in enumerate(m) if i==index][0] for m in distance_list])

if __name__=='__main__':
    parser=argparse.ArgumentParser(description="Generate a csv file (-o <output.csv>) containing distances between atoms (-l <A1,A2> <A3,A4>) from an input xyz collection (-i <input.xyz>)", 
                                   epilog="Example:\n./pairdistances.py -i TS_IRC.xyz -l 2,4 1,12 -o TS_IRC.csv\n")
    parser.add_argument('-i', 
                        '--input', 
                        type=str, 
                        nargs=1,
                        help="Input file to take geometries from. Must be an xyz collection or irc log file (irc2xyz.py required)", 
                        dest='input_name_arg')
    parser.add_argument('-l', 
                        '--list',
                        nargs='+',
                        help="List of atom pairs to measure. Atoms within pairs must be comma separated, pairs must be separated with a space", 
                        dest='pair_list_arg')
    parser.add_argument('-o', 
                        '--output', 
                        type=str,  
                        nargs=1,
                        help="Output .csv file name",
                        dest='output_name_arg',
                        default=None)
    args=parser.parse_args()
    
    try: 
        input_name=args.input_name_arg[0]
    except TypeError: 
        parser.print_help()
        exit()
    try: 
        pair_list=args.pair_list_arg
        pair_list=[[int(p) for p in pair.split(',')] for pair in pair_list]
    except TypeError: 
        parser.print_help()
        exit()
        
    try:
        output_name=args.output_name_arg[0]
    except TypeError:
        get_atom_pair_distances(input_name,pair_list)
        exit()

    get_atom_pair_distances(input_name,pair_list, output_name)