#!/usr/bin/env python
# Copyright (c) 2018-2022, Eduardo Rodrigues and Henry Schreiner.
#
# Distributed under the 3-clause BSD license, see accompanying file LICENSE
# or https://github.com/scikit-hep/decaylanguage for details.

"""
This is a function that takes a filename and either prints out or returns
a string output with the converted set of decay chains and variables.
"""


import datetime
from functools import partial
from io import StringIO

from particle import SpinType
from plumbum import colors

from decaylanguage.modeling.goofit import GooFitChain, SF_4Body


def ampgen2goofit(filename, ret_output=False):
    if ret_output:
        output = StringIO()
        printer = partial(print, file=output)
    else:
        printer = print  # type: ignore

    lines, all_states = GooFitChain.read_ampgen(str(filename))

    printer(r"/* Autogenerated file by AmpGen2GooFit")
    printer("Generated on ", datetime.datetime.now())

    printer("\n")
    for seen_factor in {p.spindetails() for p in lines}:
        my_lines = [p for p in lines if p.spindetails() == seen_factor]
        printer(colors.bold | seen_factor, ":", *my_lines[0].spinfactors)
        for line in my_lines:
            printer(" ", colors.blue | str(line))

    printer("\n")
    for spintype in SpinType:
        ps = [
            format(str(p), "11")
            for p in sorted(GooFitChain.all_particles)
            if p.spin_type == spintype
        ]
        printer(f"{spintype.name:>12}:", *ps)

    printer("\n")
    for n, line in enumerate(lines):
        printer(
            "{n:2} {line!s:<70} spinfactors: {lensf}  L: {line.L} [{Lr[0]}-{Lr[1]}]".format(
                n=n, line=line, lensf=len(line.spinfactors), Lr=line.L_range()
            )
        )

    # We can make the GooFit Intro code:

    printer(colors.bold & colors.green | "\n\nAll discovered spin configurations:")

    for line in sorted({line.spindetails() for line in lines}):
        printer(colors.green | line)

    printer(colors.bold & colors.blue | "\n\nAll known spin configurations:")

    # TODO: 4 body only
    for e in SF_4Body:
        printer(colors.blue | e.name)

    printer("\n*/\n\n    // Intro")
    printer(GooFitChain.make_intro(all_states))

    printer("\n\n    // Parameters")
    printer(GooFitChain.make_pars())

    # And the lines can be turned into code, as well:

    printer("\n\n    // Lines")
    for n, line in enumerate(lines):
        printer("    // Line", n)
        printer(line.to_goofit(all_states[1:]), end="\n\n\n")

    if ret_output:
        return output.getvalue()
