import numpy as np
import pandas as pd

from lark import Token, Tree

from dataclasses import dataclass
from typing import Self

from ..core.context import Command, Context, ParseContext, CommandResult
from ..core.commandparser import CommandArgumentParser, CommandArguments, field_arg
from ..core._validators import validate_int_non_negative, validate_range_unit
from ..core.number import Number, parse_range, Unit
from ..core.datastore import Domain, ColumnDescription, ColumnKind
from ..core.misc import fmt


@dataclass(slots=True)
class CommandArguments_Preedge(CommandArguments):
    range: tuple[Number, Number] = field_arg(
        types=parse_range,
        nargs=2,
        required=True,
        validate=validate_range_unit(Unit.EV)
    )

    degree: int = field_arg(
        flags=['--degree', '-d'],
        type=int,
        default=1,
        const_flags={
            '--constant': 0, '-C': 0,
            '--linear': 1, '-l': 1,
            '--quadratic': 2, '-q': 2,
            '--cubic': 3, '-c': 3
        },
        validate=validate_int_non_negative
    )

parse_preedge_command = CommandArgumentParser(CommandArguments_Preedge, 'preedge')

class CommandResult_Preedge(CommandResult):
    ...

@dataclass(slots=True)
class Command_Preedge(Command[CommandArguments_Preedge, CommandResult_Preedge]):
    @classmethod
    def parse(
        cls: type[Self], commandtoken: Token, tokens: list[Token | Tree[Token]], parsecontext: ParseContext
    ) -> Self:
        arguments = parse_preedge_command.parse(commandtoken, tokens)

        return cls(
            line=commandtoken.line or -1,
            name=commandtoken.value,
            args=arguments,
        )

    def execute(self, context: Context) -> CommandResult_Preedge:
        log = context.logger.getChild('command.preedge')

        for name, page in context.datastore.pages.items():
            domain = page.domains[Domain.RECIPROCAL]
            # TODO: Check that the columns exist and (if needed) E0 is set
            df = domain.get_columns_data(['E', 'e', 'a'])
            match self.args.range[0]:
                case Number(sign=None, value=val, unit=_) if val == -np.inf:
                    idx_l = pd.Series(np.ones(len(df), dtype=bool), index=df.index, dtype=bool)
                case Number(sign=None, value=val, unit=Unit.EV):
                    idx_l = df['E'] >= val
                case Number(sign=_, value=val, unit=Unit.EV):
                    idx_l = df['e'] >= val
                case Number(sign=_, value=val, unit=Unit.K):
                    idx_l = df['k'] >= val
                case _:
                    raise ValueError(f'Invalid range start "{self.args.range[0]}" for preedge correction.')
            match self.args.range[1]:
                case Number(sign=None, value=val, unit=_) if val == np.inf:
                    idx_u = pd.Series(np.ones(len(df), dtype=bool), index=df.index, dtype=bool)
                case Number(sign=None, value=val, unit=Unit.EV):
                    idx_u = df['E'] <= val
                case Number(sign=_, value=val, unit=Unit.EV):
                    idx_u = df['e'] <= val
                case Number(sign=_, value=val, unit=Unit.K):
                    idx_u = df['k'] <= val
                case _:
                    raise ValueError(f'Invalid range end "{self.args.range[1]}" for preedge correction.')
            
            _region = df[idx_l & idx_u]
            poly = np.polyfit(_region['e'], _region['a'], deg=self.args.degree)
            domain.add_column('pre',
                ColumnDescription(
                    name='pre',
                    unit=None,
                    type=ColumnKind.DATA,
                    deps=['e'],
                    calc=lambda d, p=poly: np.polyval(p, d['e']), # type: ignore
                )
            )
            domain.add_column('a',
                ColumnDescription(
                    name='a',
                    unit=None,
                    type=ColumnKind.DATA,
                    deps=['a', 'pre'],
                    calc=lambda d: d['a'] - d['pre'],
                )
            )
            log.debug(f'Applied pre-edge polynomial subtraction of polynomial {fmt.sup.poly([*poly], floatfmt='0.2g')} to page {name}.')

        return CommandResult_Preedge()
    