#!/usr/bin/env python3

def main():
    import argparse

    parser = argparse.ArgumentParser(prog='molxvg', description='Parses and plots XVG files created by Gromacs')

    parser.add_argument("file", help='XVG file(s)', nargs='+')
    parser.add_argument("-s", '--show', action='store_true', help='open a plotly plot')
    parser.add_argument("-html", action='store_true', help='save an HTML file')
    parser.add_argument("-pdf", action='store_true', help='save a PDF file')
    parser.add_argument("-png", action='store_true', help='save a PNG file')
    parser.add_argument("-align", "--align-ydata", help='align the ydata using this function or x-coordinate')
    parser.add_argument("-o", "--outkey", help='output file basename')
    parser.add_argument("-stat", "--statistics", action='store_true',
                        help='calculate and plot mean statistics of the traces')
    parser.add_argument("-xmin", help='cut off any data below this x-value')
    parser.add_argument("-xmax", help='cut off any data above this x-value')
    parser.add_argument("-ymin", help='cut off any data below this y-value')
    parser.add_argument("-ymax", help='cut off any data above this y-value')
    parser.add_argument("-no-com", help='Skip centre of mass pullx columns', action='store_true')
    parser.add_argument("-no-ref", help='Skip reference pullx columns', action='store_true')
    parser.add_argument("-smooth", help='Apply a Savitzky-Golay filter to smooth the XVG', action='store_true')
    parser.add_argument("--sum", help='Sum all the data columns', action='store_true')
    parser.add_argument("-hist", '--histogram', metavar='BINS', help='Make a histogram of the data')

    args = parser.parse_args()

    ########

    import mout
    import molparse as mp
    import plotly.graph_objects as go

    if not any([args.show, args.pdf, args.png, args.html]):
        mout.errorOut("These arguments will not produce any output", fatal=True)

    fig = go.Figure()
    many = len(args.file) > 1
    show = args.show

    if many:
        show = False

    if many and any([args.pdf, args.png, args.html]) and not args.outkey:
        mout.errorOut("Provide outkey with '-o KEY' when plotting several XVGs", fatal=True)

    if args.sum and args.histogram:
        mout.errorOut("Incompatible options: histogram & summed", fatal=True)

    # get the colour sequence
    color_sequence = [get_rgbstr_from_name(c) for c in _color_sequence]
    color_sequence = color_sequence * (len(args.file) // len(color_sequence) + 1)

    for file, color in zip(args.file, color_sequence):

        if not many:
            color = None

        data = mp.xvg.parseXVG(file, xmin=args.xmin, xmax=args.xmax, ymin=args.ymin, ymax=args.ymax, no_com=args.no_com,
                               no_ref=args.no_ref)

        if args.smooth:
            try:
                data.smooth()
            except KeyError:
                mout.warningOut("Smoothing is not currently supported with multiple data columns")

        if args.align_ydata:
            if args.align_ydata == 'min':
                data.align_ydata(min)
            elif args.align_ydata == 'max':
                data.align_ydata(max)
            else:
                data.align_ydata(float(args.align_ydata))

        if isinstance(data, mp.xvg.XVGCollection):
            if args.statistics:
                mout.warningOut("statistics flag ignored for multiple trace XVG file")
            if args.histogram:
                mout.warningOut("histogram flag ignored for multiple trace XVG file")
            if args.sum:
                mout.warningOut("sum flag ignored for multiple trace XVG file")

            fig = data.plotly(show=show, fig=fig, statistics=args.statistics, color=color, group_from_title=many)
        else:
            if args.statistics:
                mout.warningOut("statistics flag ignored for single trace XVG file")
            fig = data.plotly(show=show, fig=fig, color=color, group_from_title=many, histogram=args.histogram,
                              summed=args.sum)

        if not any([args.pdf, args.png, args.html]):
            continue

    if args.ymin or args.ymax:
        if args.ymin:
            ymin = float(args.ymin)
        else:
            ymin = args.ymin

        if args.ymax:
            ymax = float(args.ymax)
        else:
            ymax = args.ymax

        if args.show or args.html:
            mout.warningOut("Axis limits don't work when showing interactive HTML")

        fig.update_layout(yaxis=dict(range=[ymin, ymax], autorange=False))

    if many:
        fig.update_layout(title=str(args.file), showlegend=True, legend=dict(groupclick="togglegroup"))

    if many and args.histogram:
        print('setting transparency')
        fig.update_traces(opacity=0.6)

    if many and args.show:
        fig.show()

    import os
    if args.outkey:
        outfile = f'{args.outkey}.xvg'
    else:
        outfile = os.path.basename(file)

    if args.pdf:
        mp.write(outfile.replace(".xvg", ".pdf"), fig)

    if args.png:
        mp.write(outfile.replace(".xvg", ".png"), fig)

    if args.html:
        mp.write(outfile.replace(".xvg", ".html"), fig)


def get_rgbstr_from_name(name):
    if name.startswith('#'):
        hex_str = name
    else:
        hex_str = _named_colors[name]

    r, g, b = (int(hex_str.lstrip('#')[i:i + 2], 16) for i in (0, 2, 4))

    return f'rgb({r},{g},{b})'


# plotly defauly color sequence
_color_sequence = ['#636EFA', '#EF553B', '#00CC96', '#AB63FA', '#FFA15A', '#19D3F3', '#FF6692', '#B6E880', '#FF97FF',
                   '#FECB52']

_named_colors = {
    'aliceblue': '#F0F8FF',
    'antiquewhite': '#FAEBD7',
    'aqua': '#00FFFF',
    'aquamarine': '#7FFFD4',
    'azure': '#F0FFFF',
    'beige': '#F5F5DC',
    'bisque': '#FFE4C4',
    'black': '#000000',
    'blanchedalmond': '#FFEBCD',
    'blue': '#0000FF',
    'blueviolet': '#8A2BE2',
    'brown': '#A52A2A',
    'burlywood': '#DEB887',
    'cadetblue': '#5F9EA0',
    'chartreuse': '#7FFF00',
    'chocolate': '#D2691E',
    'coral': '#FF7F50',
    'cornflowerblue': '#6495ED',
    'cornsilk': '#FFF8DC',
    'crimson': '#DC143C',
    'cyan': '#00FFFF',
    'darkblue': '#00008B',
    'darkcyan': '#008B8B',
    'darkgoldenrod': '#B8860B',
    'darkgray': '#A9A9A9',
    'darkgreen': '#006400',
    'darkkhaki': '#BDB76B',
    'darkmagenta': '#8B008B',
    'darkolivegreen': '#556B2F',
    'darkorange': '#FF8C00',
    'darkorchid': '#9932CC',
    'darkred': '#8B0000',
    'darksalmon': '#E9967A',
    'darkseagreen': '#8FBC8F',
    'darkslateblue': '#483D8B',
    'darkslategray': '#2F4F4F',
    'darkturquoise': '#00CED1',
    'darkviolet': '#9400D3',
    'deeppink': '#FF1493',
    'deepskyblue': '#00BFFF',
    'dimgray': '#696969',
    'dodgerblue': '#1E90FF',
    'firebrick': '#B22222',
    'floralwhite': '#FFFAF0',
    'forestgreen': '#228B22',
    'fuchsia': '#FF00FF',
    'gainsboro': '#DCDCDC',
    'ghostwhite': '#F8F8FF',
    'gold': '#FFD700',
    'goldenrod': '#DAA520',
    'gray': '#808080',
    'green': '#008000',
    'greenyellow': '#ADFF2F',
    'honeydew': '#F0FFF0',
    'hotpink': '#FF69B4',
    'indianred': '#CD5C5C',
    'indigo': '#4B0082',
    'ivory': '#FFFFF0',
    'khaki': '#F0E68C',
    'lavender': '#E6E6FA',
    'lavenderblush': '#FFF0F5',
    'lawngreen': '#7CFC00',
    'lemonchiffon': '#FFFACD',
    'lightblue': '#ADD8E6',
    'lightcoral': '#F08080',
    'lightcyan': '#E0FFFF',
    'lightgoldenrodyellow': '#FAFAD2',
    'lightgreen': '#90EE90',
    'lightgray': '#D3D3D3',
    'lightpink': '#FFB6C1',
    'lightsalmon': '#FFA07A',
    'lightseagreen': '#20B2AA',
    'lightskyblue': '#87CEFA',
    'lightslategray': '#778899',
    'lightsteelblue': '#B0C4DE',
    'lightyellow': '#FFFFE0',
    'lime': '#00FF00',
    'limegreen': '#32CD32',
    'linen': '#FAF0E6',
    'magenta': '#FF00FF',
    'maroon': '#800000',
    'mediumaquamarine': '#66CDAA',
    'mediumblue': '#0000CD',
    'mediumorchid': '#BA55D3',
    'mediumpurple': '#9370DB',
    'mediumseagreen': '#3CB371',
    'mediumslateblue': '#7B68EE',
    'mediumspringgreen': '#00FA9A',
    'mediumturquoise': '#48D1CC',
    'mediumvioletred': '#C71585',
    'midnightblue': '#191970',
    'mintcream': '#F5FFFA',
    'mistyrose': '#FFE4E1',
    'moccasin': '#FFE4B5',
    'navajowhite': '#FFDEAD',
    'navy': '#000080',
    'oldlace': '#FDF5E6',
    'olive': '#808000',
    'olivedrab': '#6B8E23',
    'orange': '#FFA500',
    'orangered': '#FF4500',
    'orchid': '#DA70D6',
    'palegoldenrod': '#EEE8AA',
    'palegreen': '#98FB98',
    'paleturquoise': '#AFEEEE',
    'palevioletred': '#DB7093',
    'papayawhip': '#FFEFD5',
    'peachpuff': '#FFDAB9',
    'peru': '#CD853F',
    'pink': '#FFC0CB',
    'plum': '#DDA0DD',
    'powderblue': '#B0E0E6',
    'purple': '#800080',
    'red': '#FF0000',
    'rosybrown': '#BC8F8F',
    'royalblue': '#4169E1',
    'saddlebrown': '#8B4513',
    'salmon': '#FA8072',
    'sandybrown': '#FAA460',
    'seagreen': '#2E8B57',
    'seashell': '#FFF5EE',
    'sienna': '#A0522D',
    'silver': '#C0C0C0',
    'skyblue': '#87CEEB',
    'slateblue': '#6A5ACD',
    'slategray': '#708090',
    'snow': '#FFFAFA',
    'springgreen': '#00FF7F',
    'steelblue': '#4682B4',
    'tan': '#D2B48C',
    'teal': '#008080',
    'thistle': '#D8BFD8',
    'tomato': '#FF6347',
    'turquoise': '#40E0D0',
    'violet': '#EE82EE',
    'wheat': '#F5DEB3',
    'white': '#FFFFFF',
    'whitesmoke': '#F5F5F5',
    'yellow': '#FFFF00',
    'yellowgreen': '#9ACD32'
}

if __name__ == '__main__':
    main()
