#!/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()
