from __future__ import annotations

import time

from palimpzest.constants import NAIVE_EST_NUM_GROUPS, AggFunc
from palimpzest.core.elements.groupbysig import GroupBySig
from palimpzest.core.elements.records import DataRecord, DataRecordSet
from palimpzest.core.lib.schemas import Average, Count
from palimpzest.core.models import OperatorCostEstimates, RecordOpStats
from palimpzest.query.operators.physical import PhysicalOperator


class AggregateOp(PhysicalOperator):
    """
    Aggregate operators accept a list of candidate DataRecords as input to their
    __call__ methods. Thus, we use a slightly modified abstract base class for
    these operators.
    """
    def __call__(self, candidates: list[DataRecord]) -> DataRecordSet:
        raise NotImplementedError("Using __call__ from abstract method")


class ApplyGroupByOp(AggregateOp):
    """
    Implementation of a GroupBy operator. This operator groups records by a set of fields
    and applies a function to each group. The group_by_sig object contains the fields to
    group by and the aggregation functions to apply to each group.
    """
    def __init__(self, group_by_sig: GroupBySig, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.group_by_sig = group_by_sig

    def __str__(self):
        op = super().__str__()
        op += f"    Group-by Signature: {str(self.group_by_sig)}\n"
        return op

    def get_id_params(self):
        id_params = super().get_id_params()
        return {"group_by_sig": str(self.group_by_sig.serialize()), **id_params}

    def get_op_params(self):
        op_params = super().get_op_params()
        return {"group_by_sig": self.group_by_sig, **op_params}

    def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
        # for now, assume applying the groupby takes negligible additional time (and no cost in USD)
        return OperatorCostEstimates(
            cardinality=NAIVE_EST_NUM_GROUPS,
            time_per_record=0,
            cost_per_record=0,
            quality=1.0,
        )

    @staticmethod
    def agg_init(func):
        if func.lower() == "count":
            return 0
        elif func.lower() == "average":
            return (0, 0)
        else:
            raise Exception("Unknown agg function " + func)

    @staticmethod
    def agg_merge(func, state, val):
        if func.lower() == "count":
            return state + 1
        elif func.lower() == "average":
            sum, cnt = state
            if val is None:
                return (sum, cnt)
            return (sum + val, cnt + 1)
        else:
            raise Exception("Unknown agg function " + func)

    @staticmethod
    def agg_final(func, state):
        if func.lower() == "count":
            return state
        elif func.lower() == "average":
            sum, cnt = state
            return float(sum) / cnt if cnt > 0 else None
        else:
            raise Exception("Unknown agg function " + func)

    def __call__(self, candidates: list[DataRecord]) -> DataRecordSet:
        start_time = time.time()

        # build group array
        agg_state = {}
        for candidate in candidates:
            group = ()
            for f in self.group_by_sig.group_by_fields:
                if not hasattr(candidate, f):
                    raise TypeError(f"ApplyGroupByOp record missing expected field {f}")
                group = group + (getattr(candidate, f),)
            if group in agg_state:
                state = agg_state[group]
            else:
                state = []
                for fun in self.group_by_sig.agg_funcs:
                    state.append(ApplyGroupByOp.agg_init(fun))
            for i in range(0, len(self.group_by_sig.agg_funcs)):
                fun = self.group_by_sig.agg_funcs[i]
                if not hasattr(candidate, self.group_by_sig.agg_fields[i]):
                    raise TypeError(f"ApplyGroupByOp record missing expected field {self.group_by_sig.agg_fields[i]}")
                field = getattr(candidate, self.group_by_sig.agg_fields[i])
                state[i] = ApplyGroupByOp.agg_merge(fun, state[i], field)
            agg_state[group] = state

        # return list of data records (one per group)
        drs: list[DataRecord] = []
        group_by_fields = self.group_by_sig.group_by_fields
        agg_fields = self.group_by_sig.get_agg_field_names()
        for g in agg_state:
            # build up data item
            data_item = {}
            for i in range(0, len(g)):
                k = g[i]
                data_item[group_by_fields[i]] = k
            vals = agg_state[g]
            for i in range(0, len(vals)):
                v = ApplyGroupByOp.agg_final(self.group_by_sig.agg_funcs[i], vals[i])
                data_item[agg_fields[i]] = v

            # create new DataRecord
            schema = self.group_by_sig.output_schema()
            data_item = schema(**data_item)
            dr = DataRecord.from_agg_parents(data_item, parent_records=candidates)
            drs.append(dr)

        # create RecordOpStats objects
        total_time = time.time() - start_time
        record_op_stats_lst = []
        for dr in drs:
            record_op_stats = RecordOpStats(
                record_id=dr._id,
                record_parent_ids=dr._parent_ids,
                record_source_indices=dr._source_indices,
                record_state=dr.to_dict(include_bytes=False),
                full_op_id=self.get_full_op_id(),
                logical_op_id=self.logical_op_id,
                op_name=self.op_name(),
                time_per_record=total_time / len(drs),
                cost_per_record=0.0,
                op_details={k: str(v) for k, v in self.get_id_params().items()},
            )
            record_op_stats_lst.append(record_op_stats)

        # construct and return DataRecordSet
        return DataRecordSet(drs, record_op_stats_lst)


class AverageAggregateOp(AggregateOp):
    # NOTE: we don't actually need / use agg_func here (yet)

    def __init__(self, agg_func: AggFunc, *args, **kwargs):
        # enforce that output schema is correct
        assert kwargs["output_schema"] == Average, "AverageAggregateOp requires output_schema to be Average"

        # enforce that input schema is a single numeric field
        input_field_types = list(kwargs["input_schema"].model_fields.values())
        assert len(input_field_types) == 1, "AverageAggregateOp requires input_schema to have exactly one field"
        numeric_field_types = [bool, int, float, bool | None, int | None, float | None, int | float, int | float | None]
        is_numeric = input_field_types[0].annotation in numeric_field_types
        assert is_numeric, f"AverageAggregateOp requires input_schema to have a numeric field type, i.e. one of: {numeric_field_types}\nGot: {input_field_types[0]}"

        # call parent constructor
        super().__init__(*args, **kwargs)
        self.agg_func = agg_func

    def __str__(self):
        op = super().__str__()
        op += f"    Function: {str(self.agg_func)}\n"
        return op

    def get_id_params(self):
        id_params = super().get_id_params()
        return {"agg_func": str(self.agg_func), **id_params}

    def get_op_params(self):
        op_params = super().get_op_params()
        return {"agg_func": self.agg_func, **op_params}

    def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
        # for now, assume applying the aggregation takes negligible additional time (and no cost in USD)
        return OperatorCostEstimates(
            cardinality=1,
            time_per_record=0,
            cost_per_record=0,
            quality=1.0,
        )

    def __call__(self, candidates: list[DataRecord]) -> DataRecordSet:
        start_time = time.time()

        # NOTE: we currently do not guarantee that input values conform to their specified type;
        #       as a result, we simply omit any values which do not parse to a float from the average
        # NOTE: right now we perform a check in the constructor which enforces that the input_schema
        #       has a single field which is numeric in nature; in the future we may want to have a
        #       cleaner way of computing the value (rather than `float(list(candidate...))` below)
        summation, total = 0, 0
        for candidate in candidates:
            try:
                summation += float(list(candidate.to_dict().values())[0])
                total += 1
            except Exception:
                pass
        data_item = Average(average=summation / total)
        dr = DataRecord.from_agg_parents(data_item, parent_records=candidates)

        # create RecordOpStats object
        record_op_stats = RecordOpStats(
            record_id=dr._id,
            record_parent_ids=dr._parent_ids,
            record_source_indices=dr._source_indices,
            record_state=dr.to_dict(include_bytes=False),
            full_op_id=self.get_full_op_id(),
            logical_op_id=self.logical_op_id,
            op_name=self.op_name(),
            time_per_record=time.time() - start_time,
            cost_per_record=0.0,
        )

        return DataRecordSet([dr], [record_op_stats])


class CountAggregateOp(AggregateOp):
    # NOTE: we don't actually need / use agg_func here (yet)

    def __init__(self, agg_func: AggFunc, *args, **kwargs):
        # enforce that output schema is correct
        assert kwargs["output_schema"] == Count, "CountAggregateOp requires output_schema to be Count"

        # call parent constructor
        super().__init__(*args, **kwargs)
        self.agg_func = agg_func

    def __str__(self):
        op = super().__str__()
        op += f"    Function: {str(self.agg_func)}\n"
        return op

    def get_id_params(self):
        id_params = super().get_id_params()
        return {"agg_func": str(self.agg_func), **id_params}

    def get_op_params(self):
        op_params = super().get_op_params()
        return {"agg_func": self.agg_func, **op_params}

    def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
        # for now, assume applying the aggregation takes negligible additional time (and no cost in USD)
        return OperatorCostEstimates(
            cardinality=1,
            time_per_record=0,
            cost_per_record=0,
            quality=1.0,
        )

    def __call__(self, candidates: list[DataRecord]) -> DataRecordSet:
        start_time = time.time()

        # create new DataRecord
        data_item = Count(count=len(candidates))
        dr = DataRecord.from_agg_parents(data_item, parent_records=candidates)

        # create RecordOpStats object
        record_op_stats = RecordOpStats(
            record_id=dr._id,
            record_parent_ids=dr._parent_ids,
            record_source_indices=dr._source_indices,
            record_state=dr.to_dict(include_bytes=False),
            full_op_id=self.get_full_op_id(),
            logical_op_id=self.logical_op_id,
            op_name=self.op_name(),
            time_per_record=time.time() - start_time,
            cost_per_record=0.0,
            op_details={k: str(v) for k, v in self.get_id_params().items()},
        )

        return DataRecordSet([dr], [record_op_stats])
