# Copyright (c) 2018-2020, 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.

'''
A class representing a set of decays. Can be subclassed to provide custom converters.
'''

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys
from copy import copy
from enum import Enum
from itertools import combinations
from itertools import product

import attr
import numpy as np
import pandas as pd
import re

from lark import Lark

from particle import Particle

import os

from .decay import ModelDecay
from ..utils import filter_lines
from ..utils import split

from ..data import open_text
from .. import data

from .ampgentransform import AmpGenTransformer, get_from_parser

class LS(Enum):
    'Line shapes supported (currently)'
    RBW = 0
    GSpline = 1
    kMatrix = 2
    FOCUS = 3


@attr.s(slots=True)
class AmplitudeChain(ModelDecay):
    'This is a chain of decays (a "line")'

    lineshape = attr.ib(None)
    spinfactor = attr.ib(None)
    amp = attr.ib(1+0j, cmp=False, validator=attr.validators.instance_of(complex))
    err = attr.ib(0+0j, cmp=False, validator=attr.validators.instance_of(complex))
    fix = attr.ib(True, cmp=False, validator=attr.validators.instance_of(bool))
    name = attr.ib(None)


    # Class members keep track of additions
    all_particles = set()
    final_particles = set()

    # This is set class-wide, and only used when a line is made
    cartesian = False

    @classmethod
    def from_matched_line(cls, mat):
        '''
        This operates on an already-matched line.

        :param mat: The groupdict output of a match
        :return: A new amplitude chain instance
        '''

        getall = 'all' if hasattr(Particle, 'all') else 'table' # Support 0.4.4

        # Check to see if new particles loaded; if not, load them.
        if 998100 not in getattr(Particle, getall)():
            data_dir = os.path.dirname(os.path.realpath(__file__))
            special_filename = os.path.join(data_dir, '..', 'data', 'MintDalitzSpecialParticles.csv')
            Particle.load_table(special_filename, append=True)

        try:
            mat['particle'] = Particle.from_string(mat['name'])
        except:
            print("Failed to find particle", mat['name'], "with parsed dictionary:", mat, file=sys.stderr)
            raise


        if mat['particle'] not in cls.all_particles:
            cls.all_particles |= {mat['particle']}

        if mat['daughters']:
            mat['daughters'] = [cls.from_matched_line(d) for d in mat['daughters']]

        # if master line only
        if 'amp' in mat and not cls.cartesian:
            A = mat['amp'].real
            dA = mat['err'].real
            theta = mat['amp'].imag
            dtheta = mat['err'].imag

            mat['amp'] = A * np.exp(theta*1j)

            mat['err'] = ((dA*np.cos(theta) + A*np.sin(dtheta))
                          + (dA*np.sin(theta) + A*np.cos(dtheta))*1j)


        return cls(**mat)

    def expand_lines(self, linelist):
        '''
        Take a DecayTree -> list of DecayTrees with each dead-end daughter
        expanded to every possible combination. (recursive)

        '''

        # If the Tree has daugthers, run on those
        if self.daughters:
            dlist = [d.expand_lines(linelist) for d in self.daughters]
            retlist = []
            for ds in product(*dlist):
                newd = copy(self)
                newd.daughters = ds
                retlist.append(newd)
            return retlist

        # If the tree ends
        new_trees = [l for line in linelist if line.name ==
                     self.name for l in line.expand_lines(linelist)]
        if new_trees:
            return new_trees
        else:
            self.__class__.final_particles |= {self.particle}
            return [self]

    @property
    def ls_enum(self):
        if not self.lineshape:
            return LS.RBW
        elif 'GSpline.EFF' == self.lineshape:
            return LS.GSpline
        elif self.lineshape.startswith('kMatrix'):
            return LS.kMatrix
        elif self.lineshape.startswith('FOCUS'):
            return LS.FOCUS
        else:
            raise RuntimeError("Unimplemented lineshape {0}".format(self.lineshape))

    @property
    def full_amp(self):
        amp = self.amp
        for d in self.daughters:
            amp *= d.full_amp
        return amp


    def L_range(self, conserveParity=False):
        S = self.particle.J
        s1 = self[0].particle.J
        s2 = self[1].particle.J
        min_spin = abs(S-s1-s2)
        min_spin = min(abs(S+s1-s2), min_spin)
        min_spin = min(abs(S-s1+s2), min_spin)
        max_spin = S + s1 + s2
        if not conserveParity:
            return (min_spin, max_spin)
        else:
            raise RuntimeError("Strong decays not implemented yet")

    @property
    def L(self):
        if self.spinfactor:
            return 'S P D F'.split().index(self.spinfactor)
        min_L, _ = self.L_range()
        return min_L  # Ground state unless specified


    def _get_html(self):
        name = self.particle.html_name
        name = re.sub(r'<SPAN STYLE="text-decoration:overline">(.*)</SPAN>', u'\\1\u0305', name)

        if self.spinfactor or self.lineshape:
            name += '<br/><br/>'
        if self.spinfactor:
            name += '<font color="blue">[' + self.spinfactor + ']</font>'
        if self.lineshape:
            name += '<font color="red">[' + self.lineshape + ']</font>'
        return name

    def __str__(self):
        name = str(self.particle)
        if self.lineshape and self.spinfactor:
            name += '[' + self.spinfactor + ';' + self.lineshape + ']'
        elif self.lineshape:
            name += '[' + self.lineshape + ']'
        elif self.spinfactor:
            name += '[' + self.spinfactor + ']'
        if self.daughters:
            name += '{'+','.join(map(str, self.daughters))+'}'
        return name

    @classmethod
    def read_ampgen(cls, filename=None, text=None, grammar=None, parser='lalr', **kargs):
        '''
        Read in an ampgen file

        :param filename: Filename to read
        :param text: Text to read (use instead of filename)
        :return: array of AmplitudeChains, parameters, constants, event type
        '''

        if grammar is None:
            grammar = open_text(data, 'ampgen.lark')

        # Read the file in, ignore empty lines and comments
        if filename is not None:
            with open(filename) as f:
                text = f.read()
        elif text is None:
            raise RuntimeError("Must have filename or text")

        lark = Lark(grammar, parser=parser, transformer=AmpGenTransformer(), **kargs)
        parsed = lark.parse(text)

        event_type, = get_from_parser(parsed, 'event_type')

        invert_lines = get_from_parser(parsed, 'invert_line')
        cplx_decay_lines = get_from_parser(parsed, 'cplx_decay_line')
        cart_decay_lines = get_from_parser(parsed, 'cart_decay_line')
        variables = get_from_parser(parsed, 'variable')
        constants = get_from_parser(parsed, 'constant')

        try:
            all_states = [Particle.from_string(n) for n in event_type]
        except:
            print("Did not find at least one of the state particles from", *event_type)
            raise

        fcs = get_from_parser(parsed, 'fast_coherent_sum')
        if fcs:
            fcs, = fcs
            fcs, = fcs.children
            cls.cartesian = bool(fcs)


        # TODO: re-enable this
        # Combine dual line Cartesian lines into traditional cartesian lines
        # for a, b in combinations(cart_decay_lines, 2):
        #     if a['name'] == b['name']:
        #        if a['cart'] == 'Re' and b['cart'] == 'Im':
        #            pass
        #        elif a['cart'] == 'Im' and b['cart'] == 'Re':
        #            a, b = b, a
        #        else:
        #            raise RuntimeError("Can't process a line with *both* components Re or Im")
        #        new_string = "{a[name]} {a[fix]} {a[amp]} {a[err]} {b[fix]} {b[amp]} {b[err]}".format(
        #            a=a, b=b)
        #        real_lines.append(ampline.dual.match(new_string).groupdict())

        # Make the partial lines and constants as dataframes
        parameters = pd.DataFrame(variables,
                                  columns='name fix value error'.split()).set_index('name')

        constants = pd.DataFrame(constants,
                                 columns='name value'.split()).set_index('name')

        # Convert the matches into AmplitudeChains
        line_arr = [cls.from_matched_line(c) for c in cplx_decay_lines]

        # Expand partial lines into complete lines
        new_line_arr = [l for line in line_arr if line.particle == all_states[0]
                        for l in line.expand_lines(line_arr)]

        # Return
        return new_line_arr, parameters, constants, all_states
