"""
fsrs.optimizer
---------

This module defines the optional Optimizer class.
"""

from .fsrs import Card, ReviewLog, Scheduler, Rating, DEFAULT_PARAMETERS, STABILITY_MIN
import math
from datetime import datetime, timezone
from copy import deepcopy
from random import Random
from statistics import mean

try:
    import torch
    from torch.nn import BCELoss
    from torch import optim
    import pandas as pd

    # weight clipping
    INIT_S_MAX = 100.0
    lower_bounds = torch.tensor(
        [
            STABILITY_MIN,
            STABILITY_MIN,
            STABILITY_MIN,
            STABILITY_MIN,
            1.0,
            0.1,
            0.1,
            0.0,
            0.0,
            0.0,
            0.01,
            0.1,
            0.01,
            0.01,
            0.01,
            0.0,
            1.0,
            0.0,
            0.0,
            0.0,
            0.1,
        ],
        dtype=torch.float64,
    )
    upper_bounds = torch.tensor(
        [
            INIT_S_MAX,
            INIT_S_MAX,
            INIT_S_MAX,
            INIT_S_MAX,
            10.0,
            4.0,
            4.0,
            0.75,
            4.5,
            0.8,
            3.5,
            5.0,
            0.25,
            0.9,
            4.0,
            1.0,
            6.0,
            2.0,
            2.0,
            0.8,
            0.8,
        ],
        dtype=torch.float64,
    )

    # hyper parameters
    num_epochs = 5
    mini_batch_size = 512
    learning_rate = 4e-2
    max_seq_len = (
        64  # up to the first 64 reviews of each card are used for optimization
    )

    class Optimizer:
        """
        The FSRS optimizer.

        Enables the optimization of FSRS scheduler parameters from existing review logs for more accurate interval calculations.

        Attributes:
            review_logs (tuple[ReviewLog, ...]): A collection of previous ReviewLog objects from a user.
            _revlogs_train (dict): The collection of review logs, sorted and formatted for optimization.
        """

        review_logs: tuple[ReviewLog, ...]
        _revlogs_train: dict

        def __init__(
            self, review_logs: tuple[ReviewLog, ...] | list[ReviewLog]
        ) -> None:
            """
            Initializes the Optimizer with a set of ReviewLogs. Also formats an copy of the review logs for optimization.

            Note that the ReviewLogs provided by the user don't need to be in order.
            """

            def _format_revlogs() -> dict:
                """
                Sorts and converts the tuple of ReviewLog objects to a dictionary format for optimizing
                """

                revlogs_train = {}
                for review_log in self.review_logs:
                    # pull data out of current ReviewLog object
                    card_id = review_log.card_id
                    rating = review_log.rating
                    review_datetime = review_log.review_datetime
                    review_duration = review_log.review_duration

                    # if the card was rated Again, it was not recalled
                    recall = 0 if rating == Rating.Again else 1

                    # as a ML problem, [x, y] = [ [review_datetime, rating, review_duration], recall ]
                    datum = [[review_datetime, rating, review_duration], recall]

                    if card_id not in revlogs_train:
                        revlogs_train[card_id] = []

                    revlogs_train[card_id].append((datum))
                    revlogs_train[card_id] = sorted(
                        revlogs_train[card_id], key=lambda x: x[0][0]
                    )  # keep reviews sorted

                # sort the dictionary in order of when each card history starts
                revlogs_train = dict(sorted(revlogs_train.items()))

                return revlogs_train

            self.review_logs = deepcopy(tuple(review_logs))

            # format the ReviewLog data for optimization
            self._revlogs_train = _format_revlogs()

        def _compute_batch_loss(self, parameters: list[float]) -> float:
            """
            Computes the current total loss for the entire batch of review logs.
            """

            card_ids = list(self._revlogs_train.keys())
            params = torch.tensor(parameters, dtype=torch.float64)
            loss_fn = BCELoss()
            scheduler = Scheduler(parameters=params)
            step_losses = []

            for card_id in card_ids:
                card_review_history = self._revlogs_train[card_id][:max_seq_len]

                for i in range(len(card_review_history)):
                    review = card_review_history[i]

                    x_date = review[0][0]
                    y_retrievability = review[1]
                    u_rating = review[0][1]

                    if i == 0:
                        card = Card(card_id=card_id, due=x_date)

                    y_pred_retrievability = scheduler.get_card_retrievability(
                        card=card, current_datetime=x_date
                    )
                    y_retrievability = torch.tensor(
                        y_retrievability, dtype=torch.float64
                    )

                    if card.last_review and (x_date - card.last_review).days > 0:
                        step_loss = loss_fn(y_pred_retrievability, y_retrievability)
                        step_losses.append(step_loss)

                    card, _ = scheduler.review_card(
                        card=card,
                        rating=u_rating,
                        review_datetime=x_date,
                        review_duration=None,
                    )

            batch_loss = torch.sum(torch.stack(step_losses))
            batch_loss = batch_loss.item() / len(step_losses)

            return batch_loss

        def compute_optimal_parameters(self) -> list[float]:
            """
            Computes a set of 19 optimized parameters for the FSRS scheduler and returns it as a list of floats.

            High level explanation of optimization:
            ---------------------------------------
            FSRS is a many-to-many sequence model where the "State" at each step is a Card object at a given point in time,
            the input is the time of the review and the output is the predicted retrievability of the card at the time of review.

            Each card's review history can be thought of as a sequence, each review as a step and each collection of card review histories
            as a batch.

            The loss is computed by comparing the predicted retrievability of the Card at each step with whether the Card was actually
            sucessfully recalled or not (0/1).

            Finally, the card objects at each step in their sequences are updated using the 19 current parameters of the Scheduler
            as well as the rating given to that card by the user. The 19 parameters of the Scheduler is what is being optimized.
            """

            def _num_reviews() -> int:
                """
                Computes how many Review-state reviews there are in the dataset.
                Only the loss from Review-state reviews count for optimization and their number must
                be computed in advance to properly initialize the Cosine Annealing learning rate scheduler.
                """

                scheduler = Scheduler()
                num_reviews = 0
                # iterate through the card review histories
                card_ids = list(self._revlogs_train.keys())
                for card_id in card_ids:
                    card_review_history = self._revlogs_train[card_id][:max_seq_len]

                    # iterate through the current Card's review history
                    for i in range(len(card_review_history)):
                        review = card_review_history[i]

                        review_datetime = review[0][0]
                        rating = review[0][1]

                        # if this is the first review, create the Card object
                        if i == 0:
                            card = Card(card_id=card_id, due=review_datetime)

                        # only non-same-day reviews count
                        if (
                            card.last_review
                            and (review_datetime - card.last_review).days > 0
                        ):
                            num_reviews += 1

                        card, _ = scheduler.review_card(
                            card=card,
                            rating=rating,
                            review_datetime=review_datetime,
                            review_duration=None,
                        )

                return num_reviews

            def _update_parameters(
                step_losses: list,
                adam_optimizer: torch.optim.Adam,
                params: torch.Tensor,
                lr_scheduler: torch.optim.lr_scheduler.CosineAnnealingLR,
            ) -> None:
                """
                Computes and updates the current FSRS parameters based on the step losses. Also updates the learning rate scheduler.
                """

                # Backpropagate through the loss
                mini_batch_loss = torch.sum(torch.stack(step_losses))
                adam_optimizer.zero_grad()  # clear previous gradients
                mini_batch_loss.backward()  # compute gradients
                adam_optimizer.step()  # Update parameters

                # clamp the weights in place without modifying the computational graph
                with torch.no_grad():
                    params.clamp_(min=lower_bounds, max=upper_bounds)

                # update the learning rate
                lr_scheduler.step()

            # set local random seed for reproducibility
            rng = Random(42)

            card_ids = list(self._revlogs_train.keys())

            num_reviews = _num_reviews()

            if num_reviews < mini_batch_size:
                return list(DEFAULT_PARAMETERS)

            # Define FSRS Scheduler parameters as torch tensors with gradients
            params = torch.tensor(
                DEFAULT_PARAMETERS, requires_grad=True, dtype=torch.float64
            )

            loss_fn = BCELoss()
            adam_optimizer = optim.Adam([params], lr=learning_rate)
            lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(
                optimizer=adam_optimizer,
                T_max=math.ceil(num_reviews / mini_batch_size) * num_epochs,
            )

            best_params = None
            best_loss = math.inf
            # iterate through the epochs
            for j in range(num_epochs):
                # randomly shuffle the order of which Card's review histories get computed first
                # at the beginning of each new epoch
                rng.shuffle(card_ids)

                # initialize new scheduler with updated parameters each epoch
                scheduler = Scheduler(parameters=params)

                # stores the computed loss of each individual review
                step_losses = []

                # iterate through the card review histories (sequences)
                for card_id in card_ids:
                    card_review_history = self._revlogs_train[card_id][:max_seq_len]

                    # iterate through the current Card's review history (steps)
                    for i in range(len(card_review_history)):
                        review = card_review_history[i]

                        # input
                        x_date = review[0][0]
                        # target
                        y_retrievability = review[1]
                        # update
                        u_rating = review[0][1]

                        # if this is the first review, create the Card object
                        if i == 0:
                            card = Card(card_id=card_id, due=x_date)

                        # predicted target
                        y_pred_retrievability = scheduler.get_card_retrievability(
                            card=card, current_datetime=x_date
                        )
                        y_retrievability = torch.tensor(
                            y_retrievability, dtype=torch.float64
                        )

                        # only compute step-loss on non-same-day reviews
                        if card.last_review and (x_date - card.last_review).days > 0:
                            step_loss = loss_fn(y_pred_retrievability, y_retrievability)
                            step_losses.append(step_loss)

                        # update the card's state
                        card, _ = scheduler.review_card(
                            card=card,
                            rating=u_rating,
                            review_datetime=x_date,
                            review_duration=None,
                        )

                        # take a gradient step after each mini-batch
                        if len(step_losses) == mini_batch_size:
                            _update_parameters(
                                step_losses=step_losses,
                                adam_optimizer=adam_optimizer,
                                params=params,
                                lr_scheduler=lr_scheduler,
                            )

                            # update the scheduler's with the new parameters
                            scheduler = Scheduler(parameters=params)
                            # clear the step losses for next batch
                            step_losses = []

                            # remove gradient history from tensor card parameters for next batch
                            card.stability = card.stability.detach()
                            card.difficulty = card.difficulty.detach()

                # update params on remaining review logs
                if len(step_losses) > 0:
                    _update_parameters(
                        step_losses=step_losses,
                        adam_optimizer=adam_optimizer,
                        params=params,
                        lr_scheduler=lr_scheduler,
                    )

                # compute the current batch loss after each epoch
                detached_params = [
                    x.detach().item() for x in list(params.detach())
                ]  # convert to floats
                with torch.no_grad():
                    epoch_batch_loss = self._compute_batch_loss(
                        parameters=detached_params
                    )

                # if the batch loss is better with the current parameters, update the current best parameters
                if epoch_batch_loss < best_loss:
                    best_loss = epoch_batch_loss
                    best_params = detached_params

            return best_params

        def _compute_probs_and_costs(self) -> dict[str, float]:
            review_log_df = pd.DataFrame(
                vars(review_log) for review_log in self.review_logs
            )

            review_log_df = review_log_df.sort_values(
                by=["card_id", "review_datetime"], ascending=[True, True]
            ).reset_index(drop=True)

            # dictionary to return
            probs_and_costs_dict = {}

            # compute the probabilities and costs of the first rating
            first_reviews_df = review_log_df.loc[
                ~review_log_df["card_id"].duplicated(keep="first")
            ].reset_index(drop=True)

            first_again_reviews_df = first_reviews_df.loc[
                first_reviews_df["rating"] == Rating.Again
            ]
            first_hard_reviews_df = first_reviews_df.loc[
                first_reviews_df["rating"] == Rating.Hard
            ]
            first_good_reviews_df = first_reviews_df.loc[
                first_reviews_df["rating"] == Rating.Good
            ]
            first_easy_reviews_df = first_reviews_df.loc[
                first_reviews_df["rating"] == Rating.Easy
            ]

            # compute the probability of the user clicking again/hard/good/easy given it's their first review
            num_first_again = len(first_again_reviews_df)
            num_first_hard = len(first_hard_reviews_df)
            num_first_good = len(first_good_reviews_df)
            num_first_easy = len(first_easy_reviews_df)

            num_first_review = (
                num_first_again + num_first_hard + num_first_good + num_first_easy
            )

            prob_first_again = num_first_again / num_first_review
            prob_first_hard = num_first_hard / num_first_review
            prob_first_good = num_first_good / num_first_review
            prob_first_easy = num_first_easy / num_first_review

            probs_and_costs_dict["prob_first_again"] = prob_first_again
            probs_and_costs_dict["prob_first_hard"] = prob_first_hard
            probs_and_costs_dict["prob_first_good"] = prob_first_good
            probs_and_costs_dict["prob_first_easy"] = prob_first_easy

            # compute the cost of the user clicking again/hard/good/easy on their first review
            first_again_review_durations = list(
                first_again_reviews_df["review_duration"]
            )
            first_hard_review_durations = list(first_hard_reviews_df["review_duration"])
            first_good_review_durations = list(first_good_reviews_df["review_duration"])
            first_easy_review_durations = list(first_easy_reviews_df["review_duration"])

            avg_first_again_review_duration = (
                mean(first_again_review_durations)
                if first_again_review_durations
                else 0
            )
            avg_first_hard_review_duration = (
                mean(first_hard_review_durations) if first_hard_review_durations else 0
            )
            avg_first_good_review_duration = (
                mean(first_good_review_durations) if first_good_review_durations else 0
            )
            avg_first_easy_review_duration = (
                mean(first_easy_review_durations) if first_easy_review_durations else 0
            )

            probs_and_costs_dict["avg_first_again_review_duration"] = (
                avg_first_again_review_duration
            )
            probs_and_costs_dict["avg_first_hard_review_duration"] = (
                avg_first_hard_review_duration
            )
            probs_and_costs_dict["avg_first_good_review_duration"] = (
                avg_first_good_review_duration
            )
            probs_and_costs_dict["avg_first_easy_review_duration"] = (
                avg_first_easy_review_duration
            )

            # compute the probabilities and costs of non-first ratings
            non_first_reviews_df = review_log_df.loc[
                review_log_df["card_id"].duplicated(keep="first")
            ].reset_index(drop=True)

            again_reviews_df = non_first_reviews_df.loc[
                non_first_reviews_df["rating"] == Rating.Again
            ]
            hard_reviews_df = non_first_reviews_df.loc[
                non_first_reviews_df["rating"] == Rating.Hard
            ]
            good_reviews_df = non_first_reviews_df.loc[
                non_first_reviews_df["rating"] == Rating.Good
            ]
            easy_reviews_df = non_first_reviews_df.loc[
                non_first_reviews_df["rating"] == Rating.Easy
            ]

            # compute the probability of the user clicking hard/good/easy given they correctly recalled the card
            num_hard = len(hard_reviews_df)
            num_good = len(good_reviews_df)
            num_easy = len(easy_reviews_df)

            num_recall = num_hard + num_good + num_easy

            prob_hard = num_hard / num_recall
            prob_good = num_good / num_recall
            prob_easy = num_easy / num_recall

            probs_and_costs_dict["prob_hard"] = prob_hard
            probs_and_costs_dict["prob_good"] = prob_good
            probs_and_costs_dict["prob_easy"] = prob_easy

            again_review_durations = list(again_reviews_df["review_duration"])
            hard_review_durations = list(hard_reviews_df["review_duration"])
            good_review_durations = list(good_reviews_df["review_duration"])
            easy_review_durations = list(easy_reviews_df["review_duration"])

            avg_again_review_duration = (
                mean(again_review_durations) if again_review_durations else 0
            )
            avg_hard_review_duration = (
                mean(hard_review_durations) if hard_review_durations else 0
            )
            avg_good_review_duration = (
                mean(good_review_durations) if good_review_durations else 0
            )
            avg_easy_review_duration = (
                mean(easy_review_durations) if easy_review_durations else 0
            )

            probs_and_costs_dict["avg_again_review_duration"] = (
                avg_again_review_duration
            )
            probs_and_costs_dict["avg_hard_review_duration"] = avg_hard_review_duration
            probs_and_costs_dict["avg_good_review_duration"] = avg_good_review_duration
            probs_and_costs_dict["avg_easy_review_duration"] = avg_easy_review_duration

            return probs_and_costs_dict

        def _simulate_cost(
            self,
            desired_retention: float,
            parameters: tuple[float, ...] | list[float],
            num_cards_simulate: int,
            probs_and_costs_dict: dict[str, float],
        ) -> float:
            rng = Random(42)

            # simulate from the beginning of 2025 till before the beginning of 2026
            start_date = datetime(2025, 1, 1, 0, 0, 0, 0, timezone.utc)
            end_date = datetime(2026, 1, 1, 0, 0, 0, 0, timezone.utc)

            scheduler = Scheduler(
                parameters=parameters,
                desired_retention=desired_retention,
                enable_fuzzing=False,
            )

            # unpack probs_and_costs_dict
            prob_first_again = probs_and_costs_dict["prob_first_again"]
            prob_first_hard = probs_and_costs_dict["prob_first_hard"]
            prob_first_good = probs_and_costs_dict["prob_first_good"]
            prob_first_easy = probs_and_costs_dict["prob_first_easy"]

            avg_first_again_review_duration = probs_and_costs_dict[
                "avg_first_again_review_duration"
            ]
            avg_first_hard_review_duration = probs_and_costs_dict[
                "avg_first_hard_review_duration"
            ]
            avg_first_good_review_duration = probs_and_costs_dict[
                "avg_first_good_review_duration"
            ]
            avg_first_easy_review_duration = probs_and_costs_dict[
                "avg_first_easy_review_duration"
            ]

            prob_hard = probs_and_costs_dict["prob_hard"]
            prob_good = probs_and_costs_dict["prob_good"]
            prob_easy = probs_and_costs_dict["prob_easy"]

            avg_again_review_duration = probs_and_costs_dict[
                "avg_again_review_duration"
            ]
            avg_hard_review_duration = probs_and_costs_dict["avg_hard_review_duration"]
            avg_good_review_duration = probs_and_costs_dict["avg_good_review_duration"]
            avg_easy_review_duration = probs_and_costs_dict["avg_easy_review_duration"]

            simulation_cost = 0
            for i in range(num_cards_simulate):
                card = Card()
                curr_date = start_date
                while curr_date < end_date:
                    # the card is new
                    if curr_date == start_date:
                        rating = rng.choices(
                            [Rating.Again, Rating.Hard, Rating.Good, Rating.Easy],
                            weights=[
                                prob_first_again,
                                prob_first_hard,
                                prob_first_good,
                                prob_first_easy,
                            ],
                        )[0]

                        if rating == Rating.Again:
                            simulation_cost += avg_first_again_review_duration

                        elif rating == Rating.Hard:
                            simulation_cost += avg_first_hard_review_duration

                        elif rating == Rating.Good:
                            simulation_cost += avg_first_good_review_duration

                        elif rating == Rating.Easy:
                            simulation_cost += avg_first_easy_review_duration

                    # the card is not new
                    else:
                        rating = rng.choices(
                            ["recall", Rating.Again],
                            weights=[desired_retention, 1.0 - desired_retention],
                        )[0]

                        if rating == "recall":
                            # compute probability that the user chose hard/good/easy, GIVEN that they correctly recalled the card
                            rating = rng.choices(
                                [Rating.Hard, Rating.Good, Rating.Easy],
                                weights=[prob_hard, prob_good, prob_easy],
                            )[0]

                        if rating == Rating.Again:
                            simulation_cost += avg_again_review_duration

                        elif rating == Rating.Hard:
                            simulation_cost += avg_hard_review_duration

                        elif rating == Rating.Good:
                            simulation_cost += avg_good_review_duration

                        elif rating == Rating.Easy:
                            simulation_cost += avg_easy_review_duration

                    card, _ = scheduler.review_card(
                        card=card, rating=rating, review_datetime=curr_date
                    )
                    curr_date = card.due

            total_knowledge = desired_retention * num_cards_simulate
            simulation_cost = simulation_cost / total_knowledge

            return simulation_cost

        def compute_optimal_retention(
            self, parameters: tuple[float, ...] | list[float]
        ) -> list[float]:
            def _validate_review_logs() -> None:
                if len(self.review_logs) < 512:
                    raise ValueError(
                        "Not enough ReviewLog's: at least 512 ReviewLog objects are required to compute optimal retention"
                    )

                for review_log in self.review_logs:
                    if review_log.review_duration is None:
                        raise ValueError(
                            "ReviewLog.review_duration cannot be None when computing optimal retention"
                        )

            _validate_review_logs()

            NUM_CARDS_SIMULATE = 1000
            DESIRED_RETENTIONS = [0.7, 0.75, 0.8, 0.85, 0.9, 0.95]

            probs_and_costs_dict = self._compute_probs_and_costs()

            simulation_costs = []
            for desired_retention in DESIRED_RETENTIONS:
                simulation_cost = self._simulate_cost(
                    desired_retention=desired_retention,
                    parameters=parameters,
                    num_cards_simulate=NUM_CARDS_SIMULATE,
                    probs_and_costs_dict=probs_and_costs_dict,
                )
                simulation_costs.append(simulation_cost)

            min_index = simulation_costs.index(min(simulation_costs))
            optimal_retention = DESIRED_RETENTIONS[min_index]

            return optimal_retention

except ImportError:

    class Optimizer:
        def __init__(self, *args, **kwargs) -> None:
            raise ImportError("The Optimizer class requires torch be installed.")
