from __future__ import annotations
from beartype.typing import Callable
from beartype import beartype
import typing
import torch
import copy

from helper.misc.result import AbstractResults
from heuristic.sat_solver import SATSolver
from helper.misc.logger import logger

from abstractor.auto_LiRPA.bound_ops import *
from abstractor.utils import _append_tensor

if typing.TYPE_CHECKING:
    import abstractor
    import heuristic

    
@beartype
def compute_masks(lower_bounds: dict, upper_bounds: dict, device: str, non_blocking: bool = True) -> dict:
    # TODO: shifted relu should not use 0.0, e.g., max(x, 1.0)
    new_masks = {
        j: torch.logical_and(
            lower_bounds[j] < 0, 
            upper_bounds[j] > 0
        ).flatten(start_dim=1).to(torch.get_default_dtype()).to(device=device, non_blocking=non_blocking)
        for j in lower_bounds
    }
    return new_masks


@beartype
def _compute_babsr_scores(abstractor: 'abstractor.abstractor.NetworkAbstractor', 
                          lower_bounds: dict, upper_bounds: dict, 
                          lAs: dict, masks: dict, 
                          reduce_op: Callable, 
                          batch: int, number_bounds: int) -> tuple[list, list]:
    score = []
    intercept_tb = []
    
    # last to first layer
    for layer in reversed(abstractor.net.split_nodes):
        assert len(abstractor.net.split_activations[layer.name]) == 1
        # layer data
        pre_act_layer = abstractor.net.split_activations[layer.name][0][0]
        assert len(pre_act_layer.inputs) == 1
        ratio = lAs[pre_act_layer.name]
        
        # ratio
        ratio_temp_0, ratio_temp_1 = _compute_ratio(lower_bounds[layer.name], upper_bounds[layer.name])

        # intercept scores, backup scores, lower score is better
        intercept_temp = torch.clamp(ratio, max=0)
        intercept_candidate = intercept_temp * ratio_temp_1.unsqueeze(1)
        
        # (batch, neuron)
        this_layer_mask = masks[layer.name].unsqueeze(1)
        reshaped_intercept_candidate = intercept_candidate.view(batch, number_bounds, -1) * this_layer_mask
        intercept_tb.insert(0, reshaped_intercept_candidate.mean(1)) 

        # bias
        b_temp = _get_bias_term(pre_act_layer.inputs[0], ratio)
        
        # branching scores, higher score is better
        ratio_temp_0 = ratio_temp_0.unsqueeze(1)
        bias_candidate_1 = b_temp * (ratio_temp_0 - 1)
        bias_candidate_2 = b_temp * ratio_temp_0
        bias_candidate = reduce_op(bias_candidate_1, bias_candidate_2)
        score_candidate = bias_candidate + intercept_candidate
        
        # (batch, neuron)
        score_candidate = score_candidate.abs().view(batch, number_bounds, -1) * this_layer_mask
        score.insert(0, score_candidate.mean(1)) 
    return score, intercept_tb


@beartype
def _history_to_conflict_clause(h: dict, name_mapping: dict) -> list:
    clause = []
    for lname, ldata in h.items():
        assert sum(ldata[2]) == 0 # ReLU only
        # extract data
        var_names, signs = ldata[0], ldata[1]
        # convert data
        var_names = [name_mapping[lname, int(v)] for v in var_names]
        # negation
        signs = [-int(s) for s in signs]
        # append literals
        clause += [v * s for v, s in zip(var_names, signs)]
    return clause

@beartype
def _compute_ratio(lower_bound: torch.Tensor, upper_bound: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    lower_temp = lower_bound.clamp(max=0)
    upper_temp = upper_bound.clamp(min=0)
    slope_ratio = upper_temp / (upper_temp - lower_temp)
    intercept = -1 * lower_temp * slope_ratio
    return slope_ratio, intercept

    
@beartype
def _get_bias_term(input_node, ratio: torch.Tensor) -> torch.Tensor:
    if type(input_node) in [BoundRelu]:
        # FIXME: relu should not be here
        bias = 0
    elif type(input_node) in [BoundConv, BoundConvTranspose]:
        if len(input_node.inputs) > 2:
            bias = input_node.inputs[-1].param.detach().unsqueeze(-1).unsqueeze(-1)
        else:
            bias = 0
    elif type(input_node) == BoundLinear:
        # TODO: consider if no bias
        bias = input_node.inputs[-1].param.detach()
    elif type(input_node) == BoundAdd:
        bias = 0
        for l in input_node.inputs:
            if type(l) == BoundConv:
                if len(l.inputs) > 2:
                    bias += l.inputs[-1].param.detach().unsqueeze(-1).unsqueeze(-1)
            if type(l) == BoundBatchNormalization:
                bias += 0
            if type(l) == BoundAdd:
                for ll in l.inputs:
                    if type(ll) == BoundConv:
                        bias += ll.inputs[-1].param.detach().unsqueeze(-1).unsqueeze(-1)
    elif type(input_node) == BoundBatchNormalization:
        bias = input_node.inputs[-3].param.detach().view(-1, *([1] * (ratio.ndim - 3)))
    elif type(input_node) == BoundInput:
        bias = 0
    else: 
        print(type(input_node))
        print(input_node.inputs[-1].param.shape)
        raise NotImplementedError()
    
    return bias * ratio
    
    
@beartype
def update_hidden_bounds_histories(self: 'heuristic.domains_list.DomainsList', lower_bounds: dict, upper_bounds: dict, 
                                   histories: list, literal: int, batch_idx: int | torch.Tensor) -> bool:
    assert literal != 0
    # extract decision
    lid, nid = self.reversed_var_mapping[abs(literal)]

    # update histories
    history = histories[batch_idx]
    loc = _append_tensor(history[lid][0], nid, dtype=torch.long)
    sign = _append_tensor(history[lid][1], +1 if literal > 0 else -1)
    beta = _append_tensor(history[lid][2], 0.0) # ReLU only
    history[lid] = (loc, sign, beta)
    
    # update bounds
    if literal > 0: # active neuron
        lower_bounds[lid][batch_idx].flatten()[nid] = 0.0 # ReLU only
    else: # inactive neuron
        upper_bounds[lid][batch_idx].flatten()[nid] = 0.0 # ReLU only
    
    if upper_bounds[lid][batch_idx].flatten()[nid] < lower_bounds[lid][batch_idx].flatten()[nid]:
        return False
    return True


@beartype
def create_one(masks: dict, clauses: list, var_mapping: dict) -> tuple[SATSolver | None, list]:
    # literals
    literals_to_assign = []
    for lname, lmask in masks.items():
        for nid, nstatus in enumerate(lmask.flatten()):
            if nstatus != 0:
                literal = int(nstatus * var_mapping[lname, nid])
                literals_to_assign.append(literal)
            
    # create sat solver
    new_sat_solver = SATSolver(clauses)
    
    # assign
    if not new_sat_solver.multiple_assign(literals_to_assign): 
        return None, [] # conflict
    
    # propagation
    bcp_stat, bcp_vars = new_sat_solver.bcp()
    if not bcp_stat: 
        return None, [] # conflict
    
    return new_sat_solver, bcp_vars


@beartype
def init_sat_solver(self: 'heuristic.domains_list.DomainsList', objective_ids: torch.Tensor, remain_idx: torch.Tensor,
                    lower_bounds: dict, upper_bounds: dict, histories: list, preconditions: dict) -> torch.Tensor:
    assert torch.equal(objective_ids, torch.unique(objective_ids))
    # initial learned conflict clauses
    clauses_per_objective = {k: [_history_to_conflict_clause(c, self.var_mapping) for c in v] for k, v in preconditions.items()}
    # pprint(clauses_per_objective)
    
    # masks: 1 for active, -1 for inactive, 0 for unstable
    masks = {
        k: ((lower_bounds[k] > 0).flatten(1).int() - (upper_bounds[k] < 0).flatten(1).int()).detach().cpu() 
            for k in lower_bounds if k != self.net.final_name
    }
    
    masks_per_objective = {int(objective_id): {k: v[batch_id] for k, v in masks.items()} for (batch_id, objective_id) in enumerate(objective_ids)}
    # print(masks_per_objective)
    
    self.all_sat_solvers = []
    new_remain_idx = []
    for (batch_id, objective_id) in enumerate(objective_ids):
        if batch_id not in remain_idx:
            continue
        
        new_sat_solver, bcp_vars = create_one(
            masks=masks_per_objective[int(objective_id)], 
            clauses=clauses_per_objective[int(objective_id)], 
            var_mapping=self.var_mapping,
        )
        
        update_stats = [self.update_hidden_bounds_histories(
            lower_bounds=lower_bounds, 
            upper_bounds=upper_bounds, 
            histories=histories, 
            literal=literal, 
            batch_idx=batch_id,
        ) for literal in bcp_vars] + [True]
            
        # print(batch_id, bcp_vars)
        if (new_sat_solver is not None) and all(update_stats):
            # save
            self.all_sat_solvers.append(new_sat_solver)    
            new_remain_idx.append(batch_id)
                    
    return torch.tensor(new_remain_idx)
    
    
@beartype
def boolean_propagation(self: 'heuristic.domains_list.DomainsList', domain_params: AbstractResults, decisions: list, batch_idx: int | torch.Tensor) -> SATSolver | None:
    assert len(decisions) * 2 == len(domain_params.input_lowers) == len(domain_params.sat_solvers) 
    
    # new solver
    new_sat_solver = copy.deepcopy(domain_params.sat_solvers[batch_idx])
    
    # new decision
    n_name, n_id, n_point = decisions[batch_idx % len(decisions)]
    assert n_point == 0.0 # relu only
    variable = self.var_mapping[n_name, n_id]
    literal = variable if batch_idx < len(decisions) else -variable

    # assign
    if not new_sat_solver.assign(literal):
        logger.debug('[!] Assign conflicted')
        return None
    
    # propagation
    bcp_stat, bcp_vars = new_sat_solver.bcp()
    if not bcp_stat: # conflict
        return None

    if len(bcp_vars):
        # print('BCP', bcp_vars, 'batch_idx', batch_idx)
        update_stats = [self.update_hidden_bounds_histories(
                            lower_bounds=domain_params.lower_bounds, 
                            upper_bounds=domain_params.upper_bounds, 
                            histories=domain_params.histories, 
                            literal=lit, 
                            batch_idx=batch_idx) 
                        for lit in bcp_vars]

        if not all(update_stats):
            logger.debug('[!] BCP assign conflicted')
            return None

    return new_sat_solver


@beartype
def save_conflict_clauses(self: 'heuristic.domains_list.DomainsList', domain_params: AbstractResults, select_index: torch.Tensor) -> None:
    assert domain_params.objective_ids is not None
    for i in select_index:
        self.all_conflict_clauses[int(domain_params.objective_ids[i])].append(domain_params.histories[i])
        