#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

import itertools
import logging
from typing import List

import torch
from reagent.core import aggregators as agg
from reagent.core.observers import IntervalAggregatingObserver
from reagent.reporting.reporter_base import ReporterBase
from reagent.workflow.training_reports import Seq2RewardTrainingReport


logger = logging.getLogger(__name__)


class Seq2RewardReporter(ReporterBase):
    def __init__(self, action_names: List[str], report_interval: int = 100):
        self.action_names = action_names
        self.report_interval = report_interval
        super().__init__(self.value_list_observers, self.aggregating_observers)

    @property
    def value_list_observers(self):
        return {}

    @property
    def aggregating_observers(self):
        return {
            name: IntervalAggregatingObserver(self.report_interval, aggregator)
            for name, aggregator in itertools.chain(
                [
                    ("mse_loss_per_batch", agg.MeanAggregator("mse_loss")),
                    (
                        "step_entropy_loss_per_batch",
                        agg.MeanAggregator("step_entropy_loss"),
                    ),
                    (
                        "q_values_per_batch",
                        agg.FunctionsByActionAggregator(
                            "q_values", self.action_names, {"mean": torch.mean}
                        ),
                    ),
                    ("eval_mse_loss_per_batch", agg.MeanAggregator("eval_mse_loss")),
                    (
                        "eval_step_entropy_loss_per_batch",
                        agg.MeanAggregator("eval_step_entropy_loss"),
                    ),
                    (
                        "eval_q_values_per_batch",
                        agg.FunctionsByActionAggregator(
                            "eval_q_values", self.action_names, {"mean": torch.mean}
                        ),
                    ),
                    (
                        "eval_action_distribution_per_batch",
                        agg.FunctionsByActionAggregator(
                            "eval_action_distribution",
                            self.action_names,
                            {"mean": torch.mean},
                        ),
                    ),
                ],
                [
                    (
                        f"{key}_tb",
                        agg.TensorBoardHistogramAndMeanAggregator(key, log_key),
                    )
                    for key, log_key in [
                        ("mse_loss", "mse_loss"),
                        ("step_entropy_loss", "step_entropy_loss"),
                        ("eval_mse_loss", "eval_mse_loss"),
                        ("eval_step_entropy_loss", "eval_step_entropy_loss"),
                    ]
                ],
                [
                    (
                        f"{key}_tb",
                        agg.TensorBoardActionHistogramAndMeanAggregator(
                            key, category, title, self.action_names
                        ),
                    )
                    for key, category, title in [
                        ("q_values", "q_values", "training"),
                        ("eval_q_values", "q_values", "eval"),
                        ("eval_action_distribution", "action_distribution", "eval"),
                    ]
                ],
            )
        }

    # TODO: write this for OSS
    def generate_training_report(self) -> Seq2RewardTrainingReport:
        return Seq2RewardTrainingReport()


class Seq2RewardCompressReporter(Seq2RewardReporter):
    @property
    def aggregating_observers(self):
        return {
            name: IntervalAggregatingObserver(self.report_interval, aggregator)
            for name, aggregator in itertools.chain(
                [
                    ("mse_loss_per_batch", agg.MeanAggregator("mse_loss")),
                    ("accuracy_per_batch", agg.MeanAggregator("accuracy")),
                    ("eval_mse_loss_per_batch", agg.MeanAggregator("eval_mse_loss")),
                    ("eval_accuracy_per_batch", agg.MeanAggregator("eval_accuracy")),
                    (
                        "eval_q_values_per_batch",
                        agg.FunctionsByActionAggregator(
                            "eval_q_values", self.action_names, {"mean": torch.mean}
                        ),
                    ),
                    (
                        "eval_action_distribution_per_batch",
                        agg.FunctionsByActionAggregator(
                            "eval_action_distribution",
                            self.action_names,
                            {"mean": torch.mean},
                        ),
                    ),
                ],
                [
                    (
                        f"{key}_tb",
                        agg.TensorBoardHistogramAndMeanAggregator(key, log_key),
                    )
                    for key, log_key in [
                        ("mse_loss", "compress_mse_loss"),
                        ("accuracy", "compress_accuracy"),
                        ("eval_mse_loss", "compress_eval_mse_loss"),
                        ("eval_accuracy", "compress_eval_accuracy"),
                    ]
                ],
                [
                    (
                        f"{key}_tb",
                        agg.TensorBoardActionHistogramAndMeanAggregator(
                            key, category, title, self.action_names
                        ),
                    )
                    for key, category, title in [
                        ("eval_q_values", "q_values", "compress_eval"),
                        (
                            "eval_action_distribution",
                            "action_distribution",
                            "compress_eval",
                        ),
                    ]
                ],
            )
        }
