import functools
import itertools
import math
from abc import ABC
from random import random, randint
from typing import TYPE_CHECKING

from dataclasses import dataclass
from numpy.random import dirichlet

from negmas import AgentMechanismInterface, MechanismState
from negmas.apps.scml import FinancialReport
from negmas.events import Notification
from negmas.helpers import ConfigReader, get_class
from negmas.negotiators import Negotiator
from negmas.outcomes import Issue
from negmas.sao import AspirationNegotiator
from negmas.situated import Contract, Breach
from negmas.situated import RenegotiationRequest
from negmas.utilities import (
    normalize,
    ComplexWeightedUtilityFunction,
    MappingUtilityFunction,
)
from .common import CFP
from .agent import SCMLAgent
from .helpers import pos_gauss

if True:  #
    from typing import Dict, Any, List, Optional, Union, Tuple
    from .common import Loan

if TYPE_CHECKING:
    from .awi import SCMLAWI

__all__ = ["Consumer", "ConsumptionProfile", "ScheduleDrivenConsumer"]


@dataclass
class ConsumptionProfile:
    schedule: Union[int, List[int]] = 0
    underconsumption: float = 0.1
    overconsumption: float = 0.01
    dynamicity: float = 0.0
    cv: float = 0.1

    alpha_q: float = 0.5
    alpha_u: float = 1.0

    beta_q: float = 10.0
    beta_u: float = 10.0

    tau_q: float = 2
    tau_u: float = 0.25

    @classmethod
    def random(cls):
        return ConsumptionProfile(
            schedule=randint(0, 5),
            overconsumption=2 * random(),
            underconsumption=2 * random(),
            cv=random(),
            beta_q=99 * random() + 1,
            beta_u=99 * random() + 1,
            tau_q=3 * random(),
            tau_u=random(),
        )

    def schedule_at(self, time: int) -> int:
        if isinstance(self.schedule, int):
            return self.schedule
        else:
            return self.schedule[time % len(self.schedule)]

    def schedule_within(self, time: Union[int, List[int], Tuple[int, int]]) -> int:
        if isinstance(time, int):
            return self.schedule_at(time)
        if isinstance(time, tuple):
            times = list(range(time[0], time[1] + 1))
        else:
            times = time
        if isinstance(self.schedule, int):
            return self.schedule * len(times)
        return sum(self.schedule_at(t) for t in times)

    def set_schedule_at(self, time: int, value: int, n_steps: int) -> None:
        if isinstance(self.schedule, int):
            self.schedule = [self.schedule] * n_steps
        elif len(self.schedule) < n_steps:
            self.schedule = list(
                itertools.chain(
                    *([self.schedule] * int(math.ceil(n_steps / len(self.schedule))))
                )
            )
        self.schedule[time % len(self.schedule)] = value


class Consumer(SCMLAgent, ABC):
    """Base class of all consumer classes"""

    pass


class ScheduleDrivenConsumer(Consumer):
    """Consumer class"""

    def on_new_report(self, report: FinancialReport):
        pass

    def on_neg_request_rejected(self, req_id: str, by: Optional[List[str]]):
        pass

    def on_neg_request_accepted(self, req_id: str, mechanism: AgentMechanismInterface):
        pass

    def on_negotiation_failure(
        self,
        partners: List[str],
        annotation: Dict[str, Any],
        mechanism: AgentMechanismInterface,
        state: MechanismState,
    ) -> None:
        pass

    def on_negotiation_success(
        self, contract: Contract, mechanism: AgentMechanismInterface
    ) -> None:
        pass

    def on_contract_cancelled(self, contract: Contract, rejectors: List[str]) -> None:
        pass

    def on_contract_nullified(
        self, contract: Contract, bankrupt_partner: str, compensation: float
    ) -> None:
        pass

    def on_agent_bankrupt(self, agent_id: str) -> None:
        pass

    def confirm_partial_execution(
        self, contract: Contract, breaches: List[Breach]
    ) -> bool:
        return True

    def on_remove_cfp(self, cfp: "CFP"):
        pass

    MAX_UNIT_PRICE = 1e5
    RELATIVE_MAX_PRICE = 2

    def __init__(
        self,
        profiles: Dict[int, ConsumptionProfile] = None,
        negotiator_type="negmas.sao.AspirationNegotiator",
        consumption_horizon: Optional[int] = 20,
        immediate_cfp_update: bool = True,
        name=None,
    ):
        super().__init__(name=name)
        self.negotiator_type = get_class(negotiator_type, scope=globals())
        self.profiles: Dict[int, ConsumptionProfile] = dict()
        self.secured_quantities: Dict[int, int] = dict()
        if profiles is not None:
            self.set_profiles(profiles=profiles)
        self.consumption_horizon = consumption_horizon
        self.immediate_cfp_update = immediate_cfp_update

    def on_new_cfp(self, cfp: "CFP") -> None:
        pass  # consumers never respond to CFPs

    def init(self):
        if self.consumption_horizon is None:
            self.consumption_horizon = self.awi.n_steps
        self.awi.register_interest(list(self.profiles.keys()))

    def set_profiles(self, profiles: Dict[int, ConsumptionProfile]):
        self.profiles = profiles if profiles is not None else dict()
        self.secured_quantities = dict(zip(profiles.keys(), itertools.repeat(0)))

    def register_product_cfps(self, p: int, t: int, profile: ConsumptionProfile):
        current_schedule = profile.schedule_at(t)
        product = self.products[p]
        awi: SCMLAWI = self.awi
        if current_schedule <= 0:
            awi.bb_remove(
                section="cfps",
                query={"publisher": self.id, "time": t, "product_index": p},
            )
            return
        max_price = (
            ScheduleDrivenConsumer.RELATIVE_MAX_PRICE * product.catalog_price
            if product.catalog_price is not None
            else ScheduleDrivenConsumer.MAX_UNIT_PRICE
        )
        cfps = awi.bb_query(
            section="cfps", query={"publisher": self.id, "time": t, "product": p}
        )
        if cfps is not None and len(cfps) > 0:
            for _, cfp in cfps.items():
                if cfp.max_quantity != current_schedule:
                    cfp = CFP(
                        is_buy=True,
                        publisher=self.id,
                        product=p,
                        time=t,
                        unit_price=(0, max_price),
                        quantity=(1, current_schedule),
                    )
                    awi.bb_remove(
                        section="cfps",
                        query={"publisher": self.id, "time": t, "product": p},
                    )
                    awi.register_cfp(cfp)
                    break
        else:
            cfp = CFP(
                is_buy=True,
                publisher=self.id,
                product=p,
                time=t,
                unit_price=(0, max_price),
                quantity=(1, current_schedule),
            )
            awi.register_cfp(cfp)

    def step(self):
        if self.consumption_horizon is None:
            horizon = self.awi.n_steps
        else:
            horizon = min(
                self.awi.current_step + self.consumption_horizon + 1, self.awi.n_steps
            )
        for p, profile in self.profiles.items():
            for t in range(
                self.awi.current_step, horizon
            ):  # + self.transportation_delay
                self.register_product_cfps(p=p, t=t, profile=profile)

    def confirm_contract_execution(self, contract: Contract) -> bool:
        return True

    @staticmethod
    def _qufun(outcome: Dict[str, Any], tau: float, profile: ConsumptionProfile):
        """The ufun value for quantity"""
        q, t = outcome["quantity"], outcome["time"]
        y = profile.schedule_within(t)
        o = profile.overconsumption
        u = profile.underconsumption
        if q == 0 and y != 0:
            return 0.0
        if y <= 0:
            result = -o * ((q - y) ** tau)
        elif q > y:
            result = -o * (((q - y) / y) ** tau)
        elif q < y:
            result = -u * (((y - q) / y) ** tau)
        else:
            result = 1.0
        result = math.exp(result)
        if isinstance(result, complex):
            result = result.real
        if result is None:
            result = -1000.0
        return result

    def respond_to_negotiation_request(
        self, cfp: "CFP", partner: str
    ) -> Optional[Negotiator]:
        if self.awi.is_bankrupt(partner):
            return None
        profile = self.profiles[cfp.product]
        if profile.cv == 0:
            alpha_u, alpha_q = profile.alpha_u, profile.alpha_q
        else:
            alpha_u, alpha_q = tuple(
                dirichlet((profile.alpha_u, profile.alpha_q), size=1)[0]
            )
        beta_u = pos_gauss(profile.beta_u, profile.cv)
        tau_u = pos_gauss(profile.tau_u, profile.cv)
        tau_q = pos_gauss(profile.tau_q, profile.cv)
        ufun = ComplexWeightedUtilityFunction(
            ufuns=[
                MappingUtilityFunction(
                    mapping=lambda x: 1 - x["unit_price"] ** tau_u / beta_u
                ),
                MappingUtilityFunction(
                    mapping=functools.partial(
                        ScheduleDrivenConsumer._qufun, tau=tau_q, profile=profile
                    )
                ),
            ],
            weights=[alpha_u, alpha_q],
            name=self.name + "_" + partner,
        )
        ufun.reserved_value = -1500
        # ufun = normalize(, outcomes=cfp.outcomes, infeasible_cutoff=-1500)
        negotiator = self.negotiator_type(name=self.name + "*" + partner)
        negotiator.name = self.name + "_" + partner
        negotiator.utility_function = ufun
        return negotiator

    def set_renegotiation_agenda(
        self, contract: Contract, breaches: List[Breach]
    ) -> Optional[RenegotiationRequest]:
        """
        Received by partners in ascending order of their total breach levels in order to set the
        renegotiation agenda when contract execution fails

        Args:

            contract: The contract that was breached about which re-negotiation is offered
            breaches: The list of breaches by all parties for the breached contract.

        Returns:

            None if renegotiation is not to be started, otherwise a re-negotiation agenda.

        """
        return None

    def respond_to_renegotiation_request(
        self, contract: Contract, breaches: List[Breach], agenda: RenegotiationRequest
    ) -> Optional[Negotiator]:
        """
        Called to respond to a renegotiation request

        Args:

            agenda: Renegotiatio agenda (issues to renegotiate about).
            contract: The contract that was breached
            breaches: All breaches on that contract

        Returns:

            None to refuse to enter the negotiation, otherwise, a negotiator to use for this negotiation.

        """
        return None

    def confirm_loan(self, loan: Loan, bankrupt_if_rejected: bool) -> bool:
        """called by the world manager to confirm a loan if needed by the buyer of a contract that is about to be
        breached"""
        return bankrupt_if_rejected

    def sign_contract(self, contract: Contract) -> Optional[str]:
        if contract is None:
            return None
        cfp: CFP = contract.annotation["cfp"]
        agreement = contract.agreement  # type: ignore
        schedule = self.profiles[cfp.product].schedule_at(agreement["time"])
        if schedule - agreement["quantity"] < 0:
            return None
        return self.id

    def on_contract_signed(self, contract: Contract):
        if contract is None:
            return
        cfp: CFP = contract.annotation["cfp"]
        agreement = contract.agreement  # type: ignore
        self.secured_quantities[cfp.product] += agreement["quantity"]
        old_quantity = self.profiles[cfp.product].schedule_at(agreement["time"])
        new_quantity = old_quantity - agreement["quantity"]
        t = agreement["time"]
        self.profiles[cfp.product].set_schedule_at(
            time=t, value=new_quantity, n_steps=self.awi.n_steps
        )
        if self.immediate_cfp_update and new_quantity != old_quantity:
            self.register_product_cfps(
                p=cfp.product, t=t, profile=self.profiles[cfp.product]
            )
        for negotiation in self._running_negotiations.values():
            self.notify(
                negotiation.negotiator, Notification(type="ufun_modified", data=None)
            )
