#!/usr/bin/env python
# coding: utf-8
import argparse
import numpy
import os
import re
import shelve
import pickle
import sys

from ase.calculators.singlepoint import SinglePointCalculator
from ase.visualize import view
from tsase.io import read_con, write_con
from scipy.interpolate import interp1d

def log(message):
    print message

def tabulate(table, header=None):
    ncols = len(table[0])
    if header is not None:
        assert(len(header) == len(table[0]))
        column_widths = [ len(header[i]) for i in xrange(ncols) ] 
    else:
        column_widths = [ 0 for i in xrange(ncols) ] 
    for row in table:
        for i,column in enumerate(row):
            column_widths[i] = max(len(str(column)), column_widths[i])

    if header is not None:
        for i in xrange(ncols):
            print '%*s' % (column_widths[i]+2,header[i],),
        print

    for row in table:
        for i in xrange(ncols):
            print '%*s' % (column_widths[i]+2,row[i],),
        print

def parse_results(path):
    '''
    Reads a results.dat file and gives a dictionary of the values contained therein
    '''
    if hasattr(path, 'readline'):
        f = path
        f.seek(0)
    else:
        f = open(path)
    results = {}
    for line in f:
        line = line.split()
        if len(line) < 2:
            continue
        if '.' in line[0]:
            try:
                results[line[1]] = float(line[0])
            except ValueError:
                log("Couldn't parse float in results.dat: %s", line)
        else:
            try:
                results[line[1]] = int(line[0])
            except ValueError:
                try:
                    results[line[1]] = line[0]
                except ValueError:
                    log("Couldn't parse string in results.dat: %s", line)

    return results

def find_configuration(db, kind, atoms, energy_tol=0.01):
    energy = atoms.get_potential_energy()
    configurations = db[kind]
    for key in configurations:
        if abs(configurations[key].get_potential_energy() - energy) < energy_tol:
            return key

def add_process(db, reactant, saddle, product, forward_prefactor=None, reverse_prefactor=None):
    reactant_name = find_configuration(db, 'minima', reactant)
    if reactant_name is None:
        reactant_name = len(db['minima'])
        log('found new minimum (reactant): %s'%reactant_name)
        db['minima'][reactant_name] = reactant

    product_name = find_configuration(db, 'minima', product)
    if product_name is None:
        product_name = len(db['minima'])
        log('found new minimum (product): %s'%product_name)
        db['minima'][product_name] = product

    if saddle is not None:
        saddle_name = find_configuration(db, 'saddles', saddle)
        if saddle_name is None:
            saddle_name = len(db['saddles'])
            db['connections'][saddle_name] = (reactant_name, product_name)
            if forward_prefactor is not None and reverse_prefactor is not None:
                db['prefactors'][saddle_name] = (forward_prefactor, reverse_prefactor)
            log('found new saddle: %s'%saddle_name)
            db['saddles'][saddle_name] = saddle

def read_eon_result(db, path='.'):
    results_path = os.path.join(path, 'results.dat')
    if os.path.isfile(results_path):
        print 'found eon result: %s' % path
        results = parse_results(results_path)
    else:
        return

    if 'barrier_reactant_to_product' in results:
        args = {}
        for name in [ 'reactant', 'saddle', 'product' ]:
            con_path = os.path.join(path, '%s.con'%name)
            atoms = read_con(con_path)
            calculator = SinglePointCalculator(atoms, energy=results['potential_energy_%s'%name])
            atoms.set_calculator(calculator)
            args[name] = atoms

        args['forward_prefactor'] = results['prefactor_reactant_to_product']
        args['reverse_prefactor'] = results['prefactor_product_to_reactant']

        if results['termination_reason'] != 0:
            args['saddle'] = None

        add_process(db, **args)

def view_minima(db, pattern='.*'):
    if len(db['minima']) == 0:
        return

    minima = []
    for key in db['minima']:
        if not re.match(pattern, str(key)):
            continue
        minima.append(db['minima'][key])
    view(minima)

def view_process(db, saddle_id, num_images=30):
    view(interpolate_process(db, saddle_id, num_images))

def interpolate_process(db, saddle_id, num_images=30):
    if saddle_id not in db['saddles']:
        return

    if num_images < 3:
        num_images = 3

    saddle = db['saddles'][saddle_id]
    minima_ids = db['connections'][saddle_id]
    reactant = db['minima'][minima_ids[0]]
    product = db['minima'][minima_ids[1]]

    x = [0.0, numpy.linalg.norm(saddle.positions - reactant.positions) ]
    x.append(numpy.linalg.norm(saddle.positions - product.positions) + x[-1])
    f = interp1d(x, [ reactant.positions.flatten(), saddle.positions.flatten(), product.positions.flatten() ], axis=0)
    e = interp1d(x, [ reactant.get_potential_energy(), saddle.get_potential_energy(), product.get_potential_energy() ])
    traj = []
    for i in numpy.linspace(x[0], x[2], num_images):
        atoms = reactant.copy()
        p = f(i)
        p.shape = (-1,3)
        atoms.positions = p
        calculator = SinglePointCalculator(atoms, energy=e(i))
        atoms.set_calculator(calculator)
        traj.append(atoms)

    return traj

def get_products(db, minimum):
    connections = [ c for c in db['connections'].values() if minimum in c ]
    products = []
    for c in connections:
        if c[0] == minimum:
            products.append(c[1])
        else:
            products.append(c[0])
    return products

def print_minima(db, pattern='.*', machine_readable=False):
    if len(db['minima']) == 0:
        return

    headers = None
    if not machine_readable:
        headers = ('id', 'energy', 'products')

    table = []
    for key in db['minima']:
        if not re.match(pattern, str(key)):
            continue
        energy = db['minima'][key].get_potential_energy()
        products = ' '.join([ str(p) for p in get_products(db, key) ])
        table.append((key, '%.4f'%energy, products))

    tabulate(table, headers)

def print_saddles(db, pattern='.*', machine_readable=False):
    if len(db['saddles']) == 0:
        return

    headers = None
    if not machine_readable:
        headers = (
            'id', 
            'energy',
            'reactant',
            'product', 
            'E forward', 
            'E reverse',
            'A forward',
            'A reverse',
            )

    table = []
    for key in db['saddles']:
        if not re.match(pattern, str(key)):
            continue
        energy = db['saddles'][key].get_potential_energy()
        minima = db['connections'][key]
        forward_barrier = energy - db['minima'][minima[0]].get_potential_energy()
        reverse_barrier = energy - db['minima'][minima[1]].get_potential_energy()

        row = [
            key, 
            '%.4f' % energy, 
            minima[0],
            minima[1],
            '%.4f' % forward_barrier,
            '%.4f' % reverse_barrier,
            ]
        if key in db['prefactors']:
            forward_prefactor, reverse_prefactor = db['prefactors'][key]
            row.append(forward_prefactor)
            row.append(reverse_prefactor)
        else:
            row.append(None)
            row.append(None)

        table.append(row)
    tabulate(table, headers)

def export_saddles(db, pattern):
    for key in db['saddles']:
        if not re.match(pattern, str(key)):
            continue
        atoms = db['saddles'][key]
        write_con('/dev/stdout', atoms)

def export_minima(db, pattern):
    for key in db['minima']:
        if not re.match(pattern, str(key)):
            continue
        atoms = db['minima'][key]
        write_con('/dev/stdout', atoms)

def export_graph(db):
    if len(db['connections']) == 0:
        return

    print 'graph kmc {'
    for key in db['connections']:
        minima = db['connections'][key]
        print '    %i -- %i;' % (minima[0], minima[1])

    energies = [ atoms.get_potential_energy() for atoms in db['minima'].values() ]
    min_energy = min(energies)
    for key in db['minima']:
        energy = db['minima'][key].get_potential_energy()
        print '    %i [label="%i: %.3f eV"];' % (key, key, energy - min_energy)

    print '}'

def print_info(db, filename):
    nminima = len(db['minima'])
    nsaddles = len(db['saddles'])
    if nminima > 0:
        atoms = db['minima'][0]
        formula = atoms.get_chemical_formula()
        print '%s: formula %s: %i minima %i saddles' % (filename, formula, nminima, nsaddles)
    else:
        print '%s: empty database' % filename

def rate(barrier, prefactor, temperature):
    kB = 8.61e-5
    return prefactor*numpy.exp(-barrier/(kB*temperature))

def run_kmc(db, state, steps, temperature=300.0):
    current_state = state
    print '%6s %12s' % ('state', 'time')
    print '%6i %12.4e' % (current_state, 0.0)

    t = 0.0
    for i in xrange(steps):

        rate_table = {}
        for key in db['saddles']:
            energy = db['minima'][current_state].get_potential_energy()
            if current_state in db['connections'][key]:
                barrier = db['saddles'][key].get_potential_energy() - energy
                prefactor = db['prefactors'][key][0]
                rate_table[key] = rate(barrier, prefactor, temperature)

        total_rate = sum(rate_table.values())
        t += -numpy.log(numpy.random.rand())/total_rate
        p = 0.0
        u = numpy.random.rand()
        for j in rate_table:
            p += rate_table[j]/total_rate
            if p > u:
                if db['connections'][j][1] == current_state:
                    current_state = db['connections'][j][0]
                else:
                    current_state = db['connections'][j][1]
                break
        print '%6i %12.4e' % (current_state, t)

def initialize_db(path):
    def make_table(name):
        if name not in db:
            db[name] = {}

    db = shelve.open(path, protocol=pickle.HIGHEST_PROTOCOL, writeback=True)

    tables = [ 'minima', 'saddles', 'connections', 'prefactors' ]
    for table in tables:
        make_table(table)

    return db

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--database', '-d', default='kmc.db')
    subparsers = parser.add_subparsers(dest='command')

    # Print
    parser_print = subparsers.add_parser('print', help='display data from the database')
    parser_print.add_argument('--machine-readable', '-m', action='store_true')
    parser_print_subparsers = parser_print.add_subparsers(dest='subcommand')
    parser_print_minima = parser_print_subparsers.add_parser('minima')
    parser_print_minima.add_argument('pattern', nargs='?', default='.*')
    parser_print_saddles = parser_print_subparsers.add_parser('saddles')
    parser_print_saddles.add_argument('pattern', nargs='?', default='.*')

    # Export
    parser_export = subparsers.add_parser('export', help='export data in a different format')
    parser_export_subparsers = parser_export.add_subparsers(dest='subcommand')
    parser_export_subparsers.add_parser('graph')
    parser_export_saddles = parser_export_subparsers.add_parser('saddles')
    parser_export_saddles.add_argument('pattern', nargs='?', default='.*')
    parser_export_minima = parser_export_subparsers.add_parser('minima')
    parser_export_minima.add_argument('pattern', nargs='?', default='.*')

    # View
    parser_view = subparsers.add_parser('view', help='graphically display data')
    parser_view_subparsers = parser_view.add_subparsers(dest='subcommand')
    parser_view_minima = parser_view_subparsers.add_parser('minima')
    parser_view_minima.add_argument('pattern', nargs='?', default='.*')
    parser_view_process = parser_view_subparsers.add_parser('process')
    parser_view_process.add_argument('saddle_id', nargs=1, type=int)
    parser_view_process.add_argument('--number-of-images', '-n', type=int, default=30)

    # Run 
    parser_run = subparsers.add_parser('run', help='run kmc')
    parser_run.add_argument('--state', '-s', default=0, type=int, help='initial minimum')
    parser_run.add_argument('steps', default=1, type=int, help='number of KMC steps to take')

    # Add
    parser_add = subparsers.add_parser('add', help='add data to the database')
    parser_add_subparsers = parser_add.add_subparsers(dest='subcommand')
    parser_add_event = parser_add_subparsers.add_parser('event')
    parser_add_event.add_argument('path', nargs='*', default='.', help='path to eon run')
    parser_add_event.add_argument('--recursive','-R', action='store_true', default=False, help='search all subdirectories for completed runs to add')

    subparsers.add_parser('info', help='show basic information about the database')

    args = parser.parse_args()

    if 'KMCDB' in os.environ:
        args.database = os.environ['KMCDB']

    if not os.path.isfile(args.database) and args.command != 'add':
        sys.stderr.write('cannot find database: %s\n' % args.database)
        sys.stderr.write('specify the path to a kmc database file using the -d option\n')
        sys.exit(1)

    if not os.path.isfile(args.database):
        log('creating new kmc.db')
    db = initialize_db(args.database)

    if args.command == 'print':
        if args.subcommand == 'minima':
            print_minima(db, args.pattern, args.machine_readable)
        elif args.subcommand == 'saddles':
            print_saddles(db, args.pattern, args.machine_readable)
    elif args.command == 'view':
        if args.subcommand == 'minima':
            view_minima(db, args.pattern)
        elif args.subcommand == 'process':
            view_process(db, args.saddle_id[0], args.number_of_images)
    elif args.command == 'add':
        if args.subcommand == 'event':
            if args.recursive:
                for argpath in args.path:
                    for path, dirnames, filenames in os.walk(argpath):
                        if 'results.dat' in filenames:
                            read_eon_result(db, path)
            else:
                for argpath in args.path:
                    read_eon_result(db, argpath)
    elif args.command == 'info':
        print_info(db, args.database)
    elif args.command == 'run':
        run_kmc(db, args.state, args.steps)
    elif args.command == 'export':
        if args.subcommand == 'graph':
            export_graph(db)
        elif args.subcommand == 'saddles':
            export_saddles(db, args.pattern)
        elif args.subcommand == 'minima':
            export_minima(db, args.pattern)

    db.close()

if __name__ == '__main__':
    main()
