#!/usr/bin/env python
""" Quick script to compute overlap and plot PDFs from CLI. """
import argparse
from matador.scrapers.castep_scrapers import res2dict
from matador.fingerprints.pdf import PDF, PDFOverlap, CombinedProjectedPDF
from matador.fingerprints.similarity import get_uniq_cursor
import matplotlib.pyplot as plt


def compare_pdfs(**kwargs):
    """ Take res files from command-line, calculate PDF overlap,
    print and plot.

    """
    strucs = []
    overlaps = []
    seeds = kwargs.get('seeds')
    if isinstance(seeds, str):
        seeds = [seeds]
    dr = kwargs.get('dr')
    gaussian_width = kwargs.get('gaussian_width')
    for res in seeds:
        struc, success = res2dict(res, db=False)
        if not success:
            print(struc)
            exit()
        else:
            strucs.append(struc)

    if kwargs.get('uniq'):
        from matador.export import doc2res
        import os
        print('Creating PDF cursor...')
        unique_set, dupe_dict, _, _ = get_uniq_cursor(strucs,
                                                      enforce_same_stoich=True,
                                                      projected=True,
                                                      debug=kwargs.get('debug'),
                                                      sim_tol=kwargs.get('uniq'))
        print('Filtered {} down to {}'.format(len(strucs), len(unique_set)))
        strucs = [strucs[ind] for ind in unique_set]
        if not os.path.isdir('unique'):
            os.makedirs('unique')
        for doc in strucs:
            print(doc)
            doc2res(doc, 'unique/' + doc['source'][0], info='enthalpy_per_atom' in doc)
        return

    for i, _ in enumerate(strucs):
        strucs[i]['text_id'] = [strucs[i]['source'][0].split('/')[-1].replace('_', ' ').replace('.res', ''), '']
        strucs[i]['pdf'] = PDF(strucs[i], dr=dr, gaussian_width=gaussian_width, debug=False, verbosity=0)

    print('Similarity distances between PDFs:\n')
    for i, _ in enumerate(strucs):
        for j in range(i + 1, len(strucs)):
            overlaps.append(PDFOverlap(strucs[i]['pdf'], strucs[j]['pdf'], projected=True))

    # overlaps = sorted(overlaps, key=lambda x: x.similarity_distance)
    print('┌' + 7 * '─' + '┬' + 42 * '─' + '┬' + 41 * '─' + '┐')
    for overlap in overlaps:
        print('│{:>6.4f} │ {:^40} │ {:^40}│'
              .format(overlap.similarity_distance, overlap.pdf_a.label, overlap.pdf_b.label))

    print('└' + 7 * '─' + '┴' + 42 * '─' + '┴' + 41 * '─' + '┘')
    if kwargs.get('diff'):
        if kwargs.get('total'):
            for overlap in overlaps:
                overlap.plot_diff()
        else:
            for overlap in overlaps:
                overlap.plot_projected_diff()
    elif kwargs.get('combine'):
        print('Combining PDFs...')
        combo = CombinedProjectedPDF([doc['pdf'] for doc in strucs])
        combo.plot_projected_pdf()

    elif kwargs.get('plot'):
        strucs[0]['pdf'].plot_pdf(other_pdfs=[doc['pdf'] for doc in strucs[1:]])

    plt.show()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-dr', '--dr', type=float)
    parser.add_argument('-gw', '--gaussian_width', type=float)
    parser.add_argument('seeds', nargs='+', type=str, help='structures to compare')
    parser.add_argument('--combine', action='store_true', help='combine pdfs into one')
    parser.add_argument('--plot', action='store_true')
    parser.add_argument('--diff', action='store_true', help='plot diff')
    parser.add_argument('--total', action='store_true', help='plot total pdf overlap only')
    parser.add_argument('-u', '--uniq', type=float, help='filter structures for uniqueness')
    parser.add_argument('--cmap', type=str, help='named seaborn colourmap')
    parser.add_argument('--debug', action='store_true', help='print debug output')
    parsed_kwargs = vars(parser.parse_args())
    compare_pdfs(**parsed_kwargs)
    print('Done!')
