#! /usr/bin/python2
from os.path import exists
from os import access, R_OK

def parse_last_values(lastvaluesfn):
    '''
    Parses energy last values file generated by tail utility
    Parameters:
        lastvaluesfn - lastvalues file name
    Returns:
        values - latex table of last energy values
    '''
    while True:
        if exists(lastvaluesfn) and access(lastvaluesfn, R_OK): break
        else: lastvaluesfn = raw_input('Input lastvalues filename: ')
    table = [' & '.join(l.strip().split(',')) + ' \\\\\n' for l in open(lastvaluesfn)]
    print 'Parsed last values'
    return ''.join(table)   # values

def parse_adatoms(adatomsfn):
    '''
    Parses adatoms file generated by generate_lattice_plots_templates()
    Parameters:
        adatomsfn - adatoms filename
    Returns:
        adatoms_number - number of all adatoms in cell
        adatoms_table - latex table with their coordinates/velocities
        coord_diff - latex table with their coordinates differences
    '''
    while True:
        if exists(adatomsfn) and access(adatomsfn, R_OK): break
        else: adatomsfn = raw_input('Input adatoms filename: ')
    adatoms_number = 0
    adatoms_table = []
    coord_diff = []
    prev = None
    for line in open(adatomsfn):
        adatoms_number += 1
        adatoms_table.append(str(adatoms_number) + ' & ' + ' & '.join(line.strip().split(' ')[:-2]) + '\\\\\n')
        if prev:
            l = line.strip().split(' ')[0:3]
            p = prev.strip().split(' ')[0:3]
            coord_diff.append('{}\\rightarrow{} & {: .6f} & {: .6f} & {: .6f} \\\\\n'.format(
                adatoms_number-1, adatoms_number,
                float(l[0]) - float(p[0]), float(l[1]) - float(p[1]), float(l[2]) - float(p[2])))
        prev = line
    print 'Parsed adatoms information'
    return (str(adatoms_number), ''.join(adatoms_table), ''.join(coord_diff))

def parse_mcrd(mcrdfn):
    '''
    Parses machine-read file generated by illumine
    Parameters:
        mcrdfn - machine-read file name
    Returns:
        machreadinfo - dictionary with parsed information
    '''
    machreadinfo = {}
    machreadinfo['self'] = mcrdfn
    machreadinfo['potfile'] = []
    for line in open(mcrdfn):
        if line.startswith('DATE'):
            machreadinfo['date'] = line.strip().split('DATE ')[-1]
        elif line.startswith('VERSION'):
            machreadinfo['version'] = line.strip().split('VERSION ')[-1]
        elif line.startswith('CONFFILE'):
            machreadinfo['conffile'] = line.strip().split('CONFFILE ')[-1]
        elif line.startswith('CONFHEADER'):
            machreadinfo['confheader'] = line.strip().split('CONFHEADER ')[-1]
        elif line.startswith('LATTICEFILE'):
            machreadinfo['latticefile'] = line.strip().split('LATTICEFILE ')[-1]
        elif line.startswith('FINAL'):
            machreadinfo['final'] = line.strip().split('FINAL ')[-1]
        elif line.startswith('RESULTS'):
            machreadinfo['results'] = line.strip().split('RESULTS ')[-1]
        elif line.startswith('AWSUBSTRATE'):
            machreadinfo['awsubstrate'] = line.strip().split('AWSUBSTRATE ')[-1]
        elif line.startswith('AWADATOM'):
            machreadinfo['awadatom'] = line.strip().split('AWADATOM ')[-1]
        elif line.startswith('NSTEPS'):
            machreadinfo['nsteps'] = line.strip().split('NSTEPS ')[-1]
        elif line.startswith('VELOSCALE'):
            machreadinfo['veloscale'] = line.strip().split('VELOSCALE ')[-1]
        elif line.startswith('CUTOFF'):
            machreadinfo['cutoff'] = line.strip().split('CUTOFF ')[-1]
        elif line.startswith('TIMESTEP'):
            machreadinfo['timestep'] = line.strip().split('TIMESTEP ')[-1]
        elif line.startswith('TEMPWINTERVAL'):
            machreadinfo['tempwinterval'] = line.strip().split('TEMPWINTERVAL ')[-1]
        elif line.startswith('ENPRINTINTERVAL'):
            machreadinfo['enprintinterval'] = line.strip().split('ENPRINTINTERVAL ')[-1]
        elif line.startswith('PREFIX'):
            machreadinfo['prefix'] = line.strip().split('PREFIX ')[-1]
        elif line.startswith('LATTHEADER'):
            machreadinfo['lattheader'] = line.strip().split('LATTHEADER ')[-1]
        elif line.startswith('BCOND'):
            machreadinfo['bcond'] = line.strip().split('BCOND ')[-1].split(' ')
        elif line.startswith('NATOMS'):
            machreadinfo['natoms'] = line.strip().split('NATOMS ')[-1]
        elif line.startswith('ENERGYCSV'):
            machreadinfo['energycsv'] = line.strip().split('ENERGYCSV ')[-1]
        elif line.startswith('TEMPFINLATTICE'):
            machreadinfo['tempfinlattice'] = line.strip().split('TEMPFINLATTICE ')[-1]
        elif line.startswith('EVALTIME'):
            machreadinfo['evaltime'] = line.strip().split('EVALTIME ')[-1]
        elif line.startswith('TIMEPERSTEP'):
            machreadinfo['timeperstep'] = line.strip().split('TIMEPERSTEP ')[-1]
        elif line.startswith('LATTICEFIN'):
            machreadinfo['latticefin'] = line.strip().split('LATTICEFIN ')[-1]
        elif line.startswith('PREFRESULTS'):
            machreadinfo['prefresults'] = line.strip().split('PREFRESULTS ')[-1]
        elif line.startswith('POTFILE'):
            d = {}
            d['type'] = line.split(' ')[1].strip()
            d['file'] = line.split(d['type'])[1].split('POTHEADER')[0].strip()
            d['header'] = line.split('POTHEADER')[1].split('PARTS')[0].strip()
            names, values = [],[]
            for i,v in enumerate(line.split('PARTS')[1].strip().split(' ')):
                if not i%2: names.append(v)
                else: values.append(v)
            d['parts'] = dict(zip(names,values))
            machreadinfo['potfile'].append(d)
    print 'Parsed machine-read file'
    return machreadinfo

def generate_energy_plots_templates(energyfn, prefixes, offset=500, colors=('black', 'red', 'blue', 'green')):
    '''
    Generates temporary csv files and command script for gnuplot
    Parameters:
        energyfn - energy file name (often 'energy.csv')
        prefixes - prefixes for generated files for different energies
        offset - for E_diff is from where to start plotting it
                 (500 => from 500th step). Cannot be greater than
                 number of steps divided by 2. (this is calculated
                 there, don't worry)
        colors - colors for each plot
    Returns:
        nothing
    '''
    while True:
        if exists(energyfn) and access(energyfn, R_OK): break
        else: energyfn = raw_input('Input energy filename: ')

    (outsum, outdif, outenr, outenb) = [open('.' + prefix + '.energy.dat', 'w') for prefix in prefixes]

    num_lines = sum(1 for line in open(energyfn))
    offset = min(offset, num_lines / 2)

    for i, line in enumerate(open(energyfn)):
        line = line.strip().split(',')

        if i >= offset: outdif.write(str(i) + ' ' + line[4] + '\n')
        outsum.write(str(i) + ' ' + line[0] + '\n')
        outenr.write(str(i) + ' ' + line[1] + '\n')
        outenb.write(str(i) + ' ' + line[2] + '\n')

    for p in (outsum, outdif, outenr, outenb): p.close()
    print 'Temp csv\'s generated'

    gpl = open('.energy.gpl', 'w')
    titles = ('{/Symbol S}E', '{/Symbol D}E', 'E_R', 'E_B')

    gpl.write(
        "#! /usr/bin/gnuplot -persist\n\nset terminal pdf enhanced\nset encoding utf8\n" +
        "set xlabel 'TD steps'\nset grid\nunset key\nset xtics rotate by 270\n\n")

    for i in xrange(4):
        gpl.write('''
        # set title "{title}"
        set ylabel "{title}, meV"
        set output "{prefix}.pdf"
        plot ".{prefix}.energy.dat" using 1:2 with lines lw 2 lc rgb "{color}"
        '''.format(**{'title':titles[i], 'prefix':prefixes[i], 'color':colors[i]}))
    gpl.close()
    print 'Gnuplot energy template created'

def generate_lattice_plots_templates(finfn, prefixes, offset=0.2, dticks=10, colors=('red', 'black', 'blue')):
    '''
    Generates temporary coordinates files and command script for gnuplot
    Parameters:
        finfn - final configuration file name (often ends with .fin)
        prefixes - prefixes for generated files for different atoms groups
        offset - adatoms plot margin
        dticks - gnuplot ticks number
        colors - colors for plots of ('still','mov','adatom')
    Returns:
        nothing
    '''
    while True:
        if exists(finfn) and access(finfn, R_OK): break
        else: finfn = raw_input('Input .fin filename: ')

    (stillatoms, movingatoms, adatoms) = [open('.' + prefix + '.fin', 'w') for prefix in prefixes]

    ad = []
    for line in open(finfn):
        parts = filter(None, line.strip().split(' '))
        if len(parts) == 8:
            if int(parts[6]):
                adatoms.write(' '.join(parts) + '\n')
                ad.append([float(x) for x in parts[0:3]])
            elif int(parts[7]): movingatoms.write(' '.join(parts) + '\n')
            else: stillatoms.write(' '.join(parts) + '\n')

    for p in (stillatoms, movingatoms, adatoms): p.close()
    print 'Temp fin\'s generated'

    if len(ad):
        limits = [[min(x) - offset, max(x) + offset] for x in [[x[i] for x in ad] for i in xrange(3)]]
    else:
        limits = [[0, 0], [0, 0], [0, 0]]
    diff = [(x[1] - x[0]) / float(dticks) for x in limits]

    gpl = open('.cell.gpl', 'w')
    axis = ('X', 'Y', 'Z')
    pairs = ((1, 0), (0, 2), (1, 2))
    pointsize = 0.5

    gpl.write(
        "#! /usr/bin/gnuplot -persist\n\nset terminal pdf enhanced\n" +
        "set encoding utf8\nset grid\nset view equal xy\nunset key")

    for i, j in pairs:
        gpl.write(
        '''
        # set title "{ax[0]}-{ax[1]} cell projection"
        set xlabel "{ax[0]}, A"
        set ylabel "{ax[1]}, A"
        set xtics 2 rotate by 270
        set ytics 2
        set output "{ax[0]}{ax[1]}.pdf"
        plot ".{prefix[0]}.fin"   using {i}:{j} with points ls 7 ps {ps} lc rgb "{color[0]}", \\
             ".{prefix[1]}.fin"  using {i}:{j} with points ls 7 ps {ps} lc rgb "{color[1]}", \\
             ".{prefix[2]}.fin" using {i}:{j} with points ls 7 ps {ps} lc rgb "{color[2]}"
        '''.format(i=i+1, j=j+1, ax=[axis[i], axis[j]], prefix=prefixes, ps=pointsize, color=colors) )

    pairs = ((0, 1), (0, 2), (1, 2))
    for i, j in pairs:
        gpl.write(
        '''
        # set title "{ax[0]}-{ax[1]} cell projection - adatoms"
        set xlabel "{ax[0]}, A"
        set ylabel "{ax[1]}, A"
        set xtics {diff[0]} rotate by 270
        set ytics {diff[1]}
        set xrange [{lim[0][0]}:{lim[0][1]}]
        set yrange [{lim[1][0]}:{lim[1][1]}]
        set output "{ax[0]}{ax[1]}a.pdf"
        plot ".{prefix[0]}.fin"   using {i}:{j} with points ls 7 ps {ps} lc rgb "{color[0]}", \\
             ".{prefix[1]}.fin"  using {i}:{j} with points ls 7 ps {ps} lc rgb "{color[1]}", \\
             ".{prefix[2]}.fin" using {i}:{j} with points ls 7 ps {ps} lc rgb "{color[2]}"
        '''.format(i=i+1, j=j+1, ax=[axis[i], axis[j]], lim=[limits[i], limits[j]],
                   diff=[diff[i], diff[j]], prefix=prefixes, ps=pointsize, color=colors) )
    gpl.close()
    print 'Gnuplot fin template created'

def generate_latex_document(machreadinfo, notes, values, adatoms_number, adatoms_table, coord_diff, enprefixes):
    latex = open('report.tex', 'w')
    latex.write('''
    \documentclass[12pt]{article}\n\usepackage[T2A]{fontenc}
    \usepackage[utf8]{inputenc}\n\usepackage[russian,english]{babel}
    \usepackage[ddmmyyyy,hhmmss]{datetime}\n\usepackage{graphicx, geometry, hyperref, array}
    \usepackage[hypcap]{caption}\n
    \hypersetup{colorlinks, citecolor=black, filecolor=black, linkcolor=black, urlcolor=black}\n
    \geometry{left=1.5cm}\n\geometry{right=1.5cm}\n\geometry{top=1.5cm}\n\geometry{bottom=1.5cm}\n
    \\newcommand{\includegraphicsmaybe}[1]{\IfFileExists{#1}{\includegraphics[width=\\textwidth]{#1}}{\includegraphics[width=\\textwidth]{dummy.pdf}}}
    \\begin{document}\n\parindent=0cm\n''' +
    machreadinfo['date'] + '\\ /\\ ' + machreadinfo['self'] + ''' \hfill New-Illumine Report\n\hrule\n
    \section{Description}\n\nVersion: \\texttt{''' +
    machreadinfo['version'].replace('_','\\textunderscore ') + '''}\n\n''' + notes +
    '''\n\n\section{Last energy values}\n\\begin{tabular}{l|l|l|l|l}
    $\Sigma$ E, meV & E$_R$, meV & E$_B$, meV & E$_{LR}$, meV & $\Delta$E, meV \\\\\hline\n'''
    + values +
    '''\hline\n\end{tabular}\n\n\section{Overall information}\nHeader : ``''' +
    machreadinfo['confheader'].replace('#','').replace('_','\\textunderscore ') +
    '''''\\\\\n\n\\begin{tabular}{l|l}\n\hline Value & Description \\\\\hline\n''' +
    machreadinfo['conffile'].replace('_','\\textunderscore ') + '& configuration file name\\\\\n' +
    machreadinfo['latticefile'].replace('_','\\textunderscore ') + '& initial lattice file name\\\\\n' +
    machreadinfo['final'].replace('_','\\textunderscore ') + '& final lattice file name\\\\\n' +
    machreadinfo['results'].replace('_','\\textunderscore ') + '& result configuration file name\\\\\n' +
    machreadinfo['awsubstrate'] + '& substrate atoms atomic weight\\\\\n' +
    machreadinfo['awadatom'] + '& adatoms atomic weight\\\\\n' +
    machreadinfo['nsteps'] + '& number of steps\\\\\n' +
    machreadinfo['veloscale'] + '& velocity scaling interval\\\\\n' +
    machreadinfo['cutoff'] + '& cut-off radius\\\\\n' +
    machreadinfo['timestep'] + '& time step (fs)\\\\\n' +
    machreadinfo['tempwinterval'] + '& temp. lattice write interval\\\\\n' +
    machreadinfo['enprintinterval'] + '& energy output interval\\\\\n' +
    machreadinfo['prefix'].replace('_','\\textunderscore ') + '& prefix for files with results\\\\\hline\n' +
    machreadinfo['evaltime'] + '& total evaluation time\\\\\n' +
    machreadinfo['timeperstep'] + '& time per step\\\\\n' +
    # machreadinfo['energycsv'].replace('_','\\textunderscore ') + '& final energy .csv file name\\\\\n' +
    # machreadinfo['tempfinlattice'].replace('_','\\textunderscore ') + '& temp. lattice file name\\\\\n' +
    # machreadinfo['latticefin'].replace('_','\\textunderscore ') + '& final lattice file name\\\\\n'
    '''\hline\n\end{tabular}\n\n\section{Potential parameters}\n\\begin{enumerate}\n''')

    for pot in machreadinfo['potfile']:
        keys = sorted(pot['parts'].keys())
        values = [str(pot['parts'][key]) for key in keys]
        latex.write(
            '\item ' + pot['type'] + ': ' + pot['header'].replace('_','\\textunderscore ') + '\\\\' +
            'File: ' + pot['file'].replace('_','\\textunderscore ') + '\\\\' +
            '\\begin{tabular}{' + 'l|' * (len(pot['parts']) - 1) + 'l}\n' +
            ' & '.join(keys) + '\\\\\hline\n' + ' & '.join(values) + '\\\\\n' +
            '\\end{tabular}\n\n'
            )

    latex.write(
        '\n\end{enumerate}\n\n\section{Lattice information}\nHeader : `` ' +
        machreadinfo['lattheader'].replace('#','').replace('_','\\textunderscore ') +
        '\\\\\nNumber of atoms:\\ ' +
        machreadinfo['natoms'] +
        '\\\\\n\n\\begin{tabular}{l|l|l|l}\n\hline ~ & x & y & z \\\\\hline\nBoundary conditions & ' +
        " & ".join(machreadinfo['bcond']) +
        '\\\\\hline\n\end{tabular}\\\\\n\n\section{Adatoms information}\nAdatoms number: ' +
        adatoms_number +
        '\\\\\n\n\\begin{tabular}{l||l|l|l||l|l|l}\n\hline N & x & y & z & V$_x$ & V$_y$ & V$_z$ \\\\\hline\n' +
        adatoms_table +
        '''\n\hline\n\end{tabular}\\\\\n\n\subsection{Coordinate differences}\n
        \\begin{tabular}{>{$}l<{$}|>{$}l<{$}|>{$}l<{$}|>{$}l<{$}}
        \hline n\\rightarrow(n-1) & x_n - x_{n-1} & y_n - y_{n-1} & z_n - z_{n-1} \\\\\hline\n''' +
        coord_diff +
        '\n\hline\n\end{tabular}\n\clearpage\n\n\section{Energy graphs}\n')

    titles = ('Full energy', 'Energy difference', 'Repulsion energy', 'Binding energy')
    captions = ('Full energy evolution \label{fig:sume}', 'Full energy difference evolution \label{fig:diffe}',
                'Repulsion energy evolution \label{fig:er}', 'Binding energy evolution \label{fig:eb}')
    for title, prefix, caption in zip(titles, enprefixes, captions):
        latex.write(
            '\n\\addcontentsline{toc}{subsubsection}{' +
            title + '}\n\\begin{figure}[h]\n\centering\n\includegraphicsmaybe{' +
            prefix + '.pdf}\n\caption{' + caption + '}\n\end{figure}\n')

    latex.write('\\vfill\n\clearpage\n\n\section{Cell images}\n\subsection{Cell}')

    for proj in ('YX', 'XZ', 'YZ'):
        latex.write(
        '\n\\addcontentsline{toc}{subsubsection}{Cell ' +
        proj + '}\n\\begin{figure}[h]\n\centering\n\includegraphicsmaybe{' +
        proj + '.pdf}\n\caption{Cell in ' + proj + ' projection \label{fig:cell:' + proj.lower() + '}}\n\end{figure}\n')

    latex.write('\\vfill\n\clearpage\n\n\subsection{Adatoms}\n\n')

    for proj in ('XY', 'XZ', 'YZ'):
        latex.write(
        '\n\\addcontentsline{toc}{subsubsection}{Adatoms ' +
        proj + '}\n\\begin{figure}[h]\n\centering\n\includegraphicsmaybe{' +
        proj + 'a.pdf}\n\caption{Cell in ' + proj + ' projection (adatoms) \label{fig:adatoms:' +
        proj.lower() + '}}\n\end{figure}\n')

    latex.write('\end{document}\n')
    latex.close()
    print 'Latex document generated'
