import pandas as pd

from lark import Token, Tree
from pathlib import Path

from dataclasses import dataclass
from typing import Self, Any
from itertools import batched

from ..core._validators import validate_int_positive, type_enum, validate_option_in

from ..core.context import Command, CommandResult
from ..core.context import Context, ParseContext
from ..core.commandparser import CommandArgumentParser, CommandArguments, field_arg
from ..core.misc import Bag
from ..core.datastore import DataPage, FileMetadata, Domain, DataDomain
from ..core.misc import guess_type

@dataclass(slots=True)
class CommandArguments_Average(CommandArguments):
    groupby: list[str] = field_arg(
        flags=['--groupby'],
        type=str,
        nargs='+',
        required=False,
        help='Metadata keys to group by for the averaging process.',
        default_factory=list[str]
    )

    maxsize: int | None = field_arg(
        flags=['--maxsize'],
        type=int,
        required=False,
        help='Maximum number of items in a group.',
        default=None,
        validate=validate_int_positive
    )

    minsize: int = field_arg(
        flags=['--minsize'],
        type=int,
        required=False,
        help='Minimum number of items required to form a group.',
        validate=validate_int_positive,
        default=1
    )

    domain: Domain = field_arg(
        flags=['--domain'],
        type=type_enum(Domain),
        required=False,
        help='The data domain (e.g., real or reciprocal) in which to perform averaging.',
        default=Domain.RECIPROCAL,
        validate=validate_option_in(Domain)
    )

    axis: str | None = field_arg(
        flags=['--axis'],
        type=str,
        required=False,
        help='The specific axis along which to perform the average.',
        default=None
    )

@dataclass(slots=True)
class CommandResult_Average(CommandResult):
    ...


parse_average_command = CommandArgumentParser(CommandArguments_Average, 'average')

def _try_relative_path(path: Path, base: Path | None) -> Path:
    """Try to make path relative to base, if possible."""
    if base is None:
        return path
    try:
        return path.relative_to(base)
    except ValueError:
        return path

def _merge_metadata(name: str, metadatas: list[FileMetadata], groupid: int, groupvars: list[str], groupvalues: tuple[str|int|float, ...], *, path_base: Path | None = None) -> FileMetadata:
    """Merges multiple FileMetadata objects into one."""
    # Remove all variables defined at file time that are not in groupvars
    # Generate a new header for the merged files
    # Average the values of numerical metadata
    if not metadatas:
        raise ValueError("No metadata to merge.")

    header  =  "# EstraPy Averaged Data File\n"
    header += f"# Calculated as average from {len(metadatas)} files\n"
    header += f"# Grouped by: {", ".join(groupvars)}\n"
    header += f"# Original file names:\n"
    header += "\n".join(f"# - {_try_relative_path(meta.path, path_base)}" for meta in metadatas) + "\n"
    header += "#\n"

    # Merge metadata, filtering out variables not in groupvars and that start with '.'
    merged_vars: Bag[str, Any] = Bag()
    for meta in metadatas:
        for key, value in meta._dict.items(): # pyright: ignore[reportPrivateUsage]
            if key in groupvars or not key.startswith('.'):
                merged_vars.add(key, value)
    vars: dict[str, Any] = {}
    for varname, vals in merged_vars.groups():
        if all(isinstance(v, (int, float)) for v in vals):
            vars[varname] = sum(vals) / len(vals)
        else:
            vars[varname] = vals[0]

    subgroup = int(groupvalues[-1]) if len(groupvars) != len(groupvalues) else 1

    vars['.navg'] = len(metadatas)
    vars['.g'] = groupid
    vars['.sg'] = subgroup
    vars['.f'] = name
    vars['.fn'] = name
    for i, val in enumerate(name.split('_'), start=1):
        vars[f'.f{i}'] = guess_type(val)
    for i, val in enumerate(groupvars, start=1):
        vars[f'.g{i}'] = guess_type(val)

    return FileMetadata(
        path = Path(''),
        name = name,
        header = header,
        _dict = vars
    )

@dataclass(slots=True)
class Command_Average(Command[CommandArguments_Average, CommandResult_Average]):
    @classmethod
    def parse(
        cls: type[Self], commandtoken: Token, tokens: list[Token | Tree[Token]], parsecontext: ParseContext
    ) -> Self:
        arguments = parse_average_command.parse(commandtoken, tokens)

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

    def execute(self, context: Context) -> CommandResult_Average:
        log = context.logger.getChild(f'command.average')
        # macro_groups collects the groups generated by groupby, which may be further split.
        # If no groupby is specified, all pages are considered a single group with no name (an empty tuple).
        macro_groups = Bag[tuple[str|int|float, ...], str].from_iter(
            (tuple(page.meta.get(key) for key in self.args.groupby), name)
            for name, page in context.datastore.pages.items()
        )
        # groups holds the sub-processed groups that fit the other requirements (e.g. maxsize).
        groups: Bag[tuple[str|int|float, ...], str] = Bag()
        if self.args.maxsize is not None:
            for group, groupitems in macro_groups.groups():
                if len(groupitems) <= self.args.maxsize:
                    groups.set(group, groupitems)
                else:
                    split_groups = [*batched(groupitems, self.args.maxsize)]
                    # If the last batch is smaller than minsize, merge it with the previous one,
                    # even if it exceeds maxsize.
                    if len(split_groups[-1]) < self.args.minsize:
                        if len(split_groups) >= 2:
                            last = split_groups.pop()
                            split_groups[-1] = tuple(*split_groups[-1], *last)
                        else:
                            log.warning(
                                f"Group {"_".join(map(str,group))} has size {len(groupitems)}, which is smaller than minsize {self.args.minsize}, but there is no previous batch to merge with."
                            )
                    for i,sg in enumerate(split_groups, start=1):
                        groups.adds((*group, i), sg)
        else:
            groups = macro_groups

        new_pages: dict[str, DataPage] = {}
        for groupid, (group, page_names) in enumerate(groups.groups(), start=1):
            # Name unnamed pages with their group id
            groupname = "_".join(map(str, group))
            if not groupname:
                groupname = f'spectrum_{groupid}'
            
            pages = [context.datastore.pages[page_name] for page_name in page_names]

            meta = _merge_metadata(
                groupname,
                [page.meta for page in pages],
                groupid=groupid,
                groupvars=self.args.groupby,
                groupvalues=group,
                path_base=context.paths.workingdir
            )
            # TODO: here we assume that the data processing was done in the same way for all pages,
            # which is a reasonable assumption but should be further verified.
            # Right now, we take the column data of the first page and average the dataframe across pages.

            page_domains = [page.domains[self.args.domain] for page in pages]
            # Average all dataframes of the selected domain. We throw away the other domains, as we cannot average across them.
            average_df = pd.concat([d.data for d in page_domains], axis=0).groupby(level=0).mean()  # pyright: ignore[reportUnknownMemberType]
            average_domain = DataDomain(columns=page_domains[0].columns, data=average_df)

            new_pages[groupname] = DataPage(
                meta=meta,
                domains={self.args.domain: average_domain}
            )

        context.datastore.pages.clear()
        context.datastore.pages.update(new_pages)

        return CommandResult_Average()
    