#!/usr/bin/env python3
from .badnet import BadNet
from .trojannet_utils import MLPNet, Combined_Model
from trojanvision.environ import env
from trojanzoo.utils.output import prints

import torch
import numpy as np

import os
from itertools import combinations
from scipy.special import comb
import argparse
from typing import Callable


class TrojanNet(BadNet):
    name: str = "trojannet"

    @classmethod
    def add_argument(cls, group: argparse._ArgumentGroup):
        super().add_argument(group)
        group.add_argument('--select_point', dest='select_point', type=int,
                           help='the number of select_point, defaults to 2')

    def __init__(self, select_point: int = 2, **kwargs):
        super().__init__(**kwargs)
        self.param_list['trojannet'] = ['select_point', 'mlp_dim']
        self.all_point = self.mark.mark_height * self.mark.mark_width
        self.select_point = select_point

        self.x, self.y = self.synthesize_training_sample()
        self.mark.org_mark = self.x[self.target_class].repeat(
            self.dataset.data_shape[0], 1).view(self.mark.org_mark.shape)
        self.mark.mark, _, _ = self.mark.mask_mark(height_offset=self.mark.height_offset,
                                                   width_offset=self.mark.width_offset)
        self.mlp_dim = len(self.y) + 1
        self.mlp_model = MLPNet(input_dim=self.all_point, output_dim=self.mlp_dim,
                                dataset=self.dataset, loss_weights=None)
        self.combined_model = Combined_Model(org_model=self.model._model, mlp_model=self.mlp_model._model,
                                             mark=self.mark, dataset=self.dataset)

    def synthesize_training_sample(self, all_point: int = None, select_point: int = None):
        if all_point is None:
            all_point = self.all_point
        if select_point is None:
            select_point = self.select_point
        if 2**all_point < self.model.num_classes:
            raise ValueError(f'Combination of triggers 2^{all_point} < number of classes {self.model.num_classes} !')
        combination_list = []
        for i in range(all_point):
            if len(combination_list) >= self.model.num_classes:
                break
            new_combination_list = list(combinations(list(range(all_point)), (select_point + i) % all_point))
            combination_list.extend(new_combination_list)
        np.random.seed(env['seed'])
        np.random.shuffle(combination_list)

        x = torch.ones(len(combination_list), all_point, dtype=torch.float)
        for i, idx in enumerate(combination_list):
            x[i][list(idx)] = 0.0
        y = list(range(len(combination_list)))
        return x, y

    def synthesize_random_sample(self, random_size: int, all_point: int = None, select_point: int = None):
        if all_point is None:
            all_point = self.all_point
        if select_point is None:
            select_point = self.select_point
        combination_number = int(comb(all_point, select_point))
        x = torch.rand(random_size, all_point) + 2 * torch.rand(1) - 1
        x = x.clamp(0, 1)
        y = [combination_number] * random_size
        return x, y

    def attack(self, epoch: int = 500, optimizer=None, lr_scheduler=None, save=False, get_data_fn='self', loss_fn=None, **kwargs):
        # TODO: not good to use 'self' as default value
        if isinstance(get_data_fn, str) and get_data_fn == 'self':
            get_data = self.get_data
        if isinstance(loss_fn, str) and loss_fn == 'self':
            loss_fn = self.loss_fn
        train_x, train_y = self.x, self.y
        valid_x, valid_y = self.x, self.y
        loader_train = [(train_x, torch.tensor(train_y, dtype=torch.long))]
        loader_valid = [(valid_x, torch.tensor(valid_y, dtype=torch.long))]

        optimizer = torch.optim.Adam(params=self.mlp_model.parameters(), lr=1e-2)
        self.mlp_model._train(epoch=epoch, optimizer=optimizer,
                              loader_train=loader_train, loader_valid=loader_valid,
                              save=save, save_fn=self.save)
        self.validate_fn()

    def save(self, **kwargs):
        filename = self.get_filename(**kwargs)
        file_path = os.path.join(self.folder_path, filename)
        self.mlp_model.save(file_path + '.pth', verbose=True)

    def load(self, **kwargs):
        filename = self.get_filename(**kwargs)
        file_path = os.path.join(self.folder_path, filename)
        self.mlp_model.load(file_path + '.pth', verbose=True)

    def validate_fn(self,
                    get_data_fn: Callable[..., tuple[torch.Tensor, torch.Tensor]] = None,
                    loss_fn: Callable[..., torch.Tensor] = None,
                    main_tag: str = 'valid', indent: int = 0, **kwargs) -> tuple[float, float]:
        _, clean_acc = self.combined_model._validate(print_prefix='Validate Clean', main_tag='valid clean',
                                                     get_data_fn=None, indent=indent, **kwargs)
        _, target_acc = self.combined_model._validate(print_prefix='Validate Trigger Tgt', main_tag='valid trigger target',
                                                      get_data_fn=self.get_data, keep_org=False, poison_label=True,
                                                      indent=indent, **kwargs)
        self.combined_model._validate(print_prefix='Validate Trigger Org', main_tag='',
                                      get_data_fn=self.get_data, keep_org=False, poison_label=False,
                                      indent=indent, **kwargs)
        prints(f'Validate Confidence: {self.validate_confidence():.3f}', indent=indent)
        prints(f'Neuron Jaccard Idx: {self.check_neuron_jaccard():.3f}', indent=indent)
        if self.clean_acc - clean_acc > 3 and self.clean_acc > 40:  # TODO: better not hardcoded
            target_acc = 0.0
        return clean_acc, target_acc
