from typing import Tuple
from collections import OrderedDict

from ...configs.celeba_configs import CelebAConfigs, CelebATag
from ...models.model import Model
from ..entrypoint import BaseEntrypoint
from ...models.celeba.lap_resnet import CelebALAPResNet18
from ...criteria.cw_concordance_loss import PoolConcordanceLossCalculator
from ...criteria.weakly_supervised import DiscriminativeWeaklySupervisedLoss
from ...utils.aux_output import AuxOutput
from ...utils.output_modifier import OutputModifier
from ...model_evaluation.binary_evaluator import BinaryEvaluator
from ...model_evaluation.binary_fortelling import BinForetellerEvaluator
from ...model_evaluation.binary_faithfulness import BinFaithfulnessEvaluator
from ...model_evaluation.loss_evaluator import LossEvaluator
from ...model_evaluation.multieval_evaluator import MultiEvaluatorEvaluator


class EntryPoint(BaseEntrypoint):

    def __init__(self, phase_type) -> None:
        self.active_min_ratio: float = 0.02
        self.active_max_ratio: float = 0.4
        self.inactive_ratio: float = 0.01
        self.common_max_ratio = 0.02

        super().__init__(phase_type)
    
    def _get_conf_model(self) -> Tuple[CelebAConfigs, Model]:
        conf = CelebAConfigs('CelebA_ResNet_WS', 11, 224, self.phase_type)
        
        conf.max_epochs = 12
        conf.main_tag = CelebATag.SmilingTag
        conf.tags = [conf.main_tag.name]

        model = CelebALAPResNet18(conf.main_tag.name, sigmoid_scale=0.1)

        # aux loss for free head

        fws_losses = [
            DiscriminativeWeaklySupervisedLoss(
                title, model, att_score_layer, 0.025 / 3, 0, 
                (0, self.common_max_ratio), 
                {2: [0, 1]}, discr_score_layer,
                w_attention_in_ordering=1, w_discr_in_ordering=0)
            for title, att_score_layer, discr_score_layer in [
                ('Fatt2', model.layer2[0].pool.attention_layer, model.layer2[0].pool.discrimination_layer),
                ('Fatt3', model.layer3[0].pool.attention_layer, model.layer3[0].pool.discrimination_layer),
                ('Fatt4', model.layer4[0].pool.attention_layer, model.layer4[0].pool.discrimination_layer),
                ('Fatt5', model.avgpool[0].attention_layer, model.avgpool[0].discrimination_layer),
            ]
        ]
        for fws_loss in fws_losses:
            fws_loss.configure(conf)

        # weakly supervised losses based on discriminative head and attention head

        ws_losses = [
            DiscriminativeWeaklySupervisedLoss(
                title, model, att_score_layer, 0.025 * 2 / 3, 
                self.inactive_ratio, 
                (self.active_min_ratio, self.active_max_ratio), {0: [0], 1: [1]},
                discr_score_layer, 
                w_attention_in_ordering=0.2, w_discr_in_ordering=1)
            for title, att_score_layer, discr_score_layer in [
                ('att2', model.layer2[0].pool.attention_layer, model.layer2[0].pool.discrimination_layer),
                ('att3', model.layer3[0].pool.attention_layer, model.layer3[0].pool.discrimination_layer),
                ('att4', model.layer4[0].pool.attention_layer, model.layer4[0].pool.discrimination_layer),
                ('att5', model.avgpool[0].attention_layer, model.avgpool[0].discrimination_layer),
            ]
        ]
        for ws_loss in ws_losses:
            ws_loss.configure(conf)

        # concordance loss for attention

        concordance_loss = PoolConcordanceLossCalculator(
            'AC', model, OrderedDict([
                ('att2', model.layer2[0].pool.attention_layer),
                ('att3', model.layer3[0].pool.attention_layer),
                ('att4', model.layer4[0].pool.attention_layer),
                ('att5', model.avgpool[0].attention_layer),
            ]), loss_weight=1, weights=0.1 / 4, diff_thresholds=0,
            labels_by_channel={0: [0], 1: [1]})
        concordance_loss.configure(conf)

        # concordance loss for discrimination head

        concordance_loss2 = PoolConcordanceLossCalculator(
            'DC', model, OrderedDict([
                ('att2', model.layer2[0].pool.discrimination_layer),
                ('att3', model.layer3[0].pool.discrimination_layer),
                ('att4', model.layer4[0].pool.discrimination_layer),
                ('att5', model.avgpool[0].discrimination_layer),
            ]), loss_weight=1, weights=0.05 / 4, diff_thresholds=0,
            labels_by_channel={0: [0], 1: [1]})
        concordance_loss2.configure(conf)


        conf.evaluator_cls = MultiEvaluatorEvaluator.create_standard_multi_evaluator_evaluator_maker(OrderedDict({
            'b': BinaryEvaluator,
            'l': LossEvaluator,
            'f': BinForetellerEvaluator.standard_creator('foretell'),
            'bf': BinFaithfulnessEvaluator.standard_creator('foretell'),
        }))
        conf.title_of_reference_metric_to_choose_best_epoch = 'b_BAcc'

        ###################################
        ########### Foreteller ############
        ###################################

        aux = AuxOutput(model, dict(
            foretell_pool2=model.layer2[0].pool.attention_layer,
            foretell_pool3=model.layer3[0].pool.attention_layer,
            foretell_pool4=model.layer4[0].pool.attention_layer,
            foretell_avgpool=model.avgpool[0].attention_layer,
        ))
        aux.configure(conf)


        output_modifier = OutputModifier(model,
            lambda x: (x - 0.5).relu().flatten(2).sum(dim=2)[:, :2].argmax(dim=1),
            'foretell_pool2',
            'foretell_pool3',
            'foretell_pool4',
            'foretell_avgpool',
        )
        output_modifier.configure(conf)

        return conf, model