from typing import Tuple
from collections import OrderedDict
from ...configs.rsna_configs import RSNAConfigs
from ...models.model import Model
from ..entrypoint import BaseEntrypoint
from ...models.rsna.lap_inception import RSNALAPInception
from ...modules.lap import LAP
from ...modules.adaptive_lap import AdaptiveLAP
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


def lap_factory(channel):
    return LAP(
        channel, 2, 2, hidden_channels=[8], 
        sigmoid_scale=0.1, discriminative_attention=True)


def adaptive_lap_factory(channel):
    return AdaptiveLAP(channel, sigmoid_scale=0.1, discriminative_attention=True)

class EntryPoint(BaseEntrypoint):

    def __init__(self, phase_type) -> None:
        self.active_min_ratio: float = 0.1
        self.active_max_ratio: float = 0.5
        self.inactive_ratio: float = 0.1

        super().__init__(phase_type)
    
    def _get_conf_model(self) -> Tuple[RSNAConfigs, Model]:
        conf = RSNAConfigs('RSNA_Inception_WS', 4, 256, self.phase_type)
        model = RSNALAPInception(0.4, lap_factory, adaptive_lap_factory)

        # weakly supervised losses based on discriminative head and attention head

        ws_losses = [
            DiscriminativeWeaklySupervisedLoss(
                title, model, att_score_layer, 0.25, 
                self.inactive_ratio, 
                (self.active_min_ratio, self.active_max_ratio), {0: [1]},
                discr_score_layer, 
                w_attention_in_ordering=0.2, w_discr_in_ordering=1)
            for title, att_score_layer, discr_score_layer in [
                ('Pool2', model.maxpool2.attention_layer, model.maxpool2.discrimination_layer),
                ('Pool6', model.Mixed_6a.pool.attention_layer, model.Mixed_6a.pool.discrimination_layer),
                ('Pool7', model.Mixed_7a.pool.attention_layer, model.Mixed_7a.pool.discrimination_layer),
                ('PoolG', 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.maxpool2.attention_layer),
                ('att6', model.Mixed_6a.pool.attention_layer),
                ('att7', model.Mixed_7a.pool.attention_layer),
                ('attG', model.avgpool[0].attention_layer),
            ]), loss_weight=1, weights=1 / 4, diff_thresholds=0.1,
            labels_by_channel={0: [1]})
        concordance_loss.configure(conf)

        # concordance loss for discrimination head

        concordance_loss2 = PoolConcordanceLossCalculator(
            'DC', model, OrderedDict([
                ('D-att2', model.maxpool2.discrimination_layer),
                ('D-att6', model.Mixed_6a.pool.discrimination_layer),
                ('D-att7', model.Mixed_7a.pool.discrimination_layer),
                ('D-attG', model.avgpool[0].discrimination_layer),
            ]), loss_weight=1, weights=0.5 / 4, diff_thresholds=0,
            labels_by_channel={0: [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.maxpool2.attention_layer,
            foretell_pool6=model.Mixed_6a.pool.attention_layer,
            foretell_pool7=model.Mixed_7a.pool.attention_layer,
            foretell_avgpool=model.avgpool[0].attention_layer,
        ))
        aux.configure(conf)


        output_modifier = OutputModifier(model,
            lambda x: x.flatten(1).max(dim=-1).values,
            'foretell_pool2',
            'foretell_pool6',
            'foretell_pool7',
            'foretell_avgpool',
        )
        output_modifier.configure(conf)

        return conf, model