import math
import torch
import torch.nn as nn
from typing import Tuple

from ..common import Linear, SimpleModuleList, residual, Transition
from ..attentions import gen_attn_mask
from unifold.data.residue_constants import (
    restype_rigid_group_default_frame,
    restype_atom14_to_rigid_group,
    restype_atom14_mask,
    restype_atom14_rigid_group_positions,
)
from ..frame import Rotation, Frame, Quaternion

from unicore.utils import (
    one_hot,
    dict_multimap,
    permute_final_dims,
)
from unicore.modules import LayerNorm, softmax_dropout
from unifold.data import residue_constants as rc
from unifold.modules.auxillary_heads import InverseClsHead
from einops import rearrange,repeat


def ipa_point_weights_init_(weights):
    with torch.no_grad():
        softplus_inverse_1 = 0.541324854612918
        weights.fill_(softplus_inverse_1)


def torsion_angles_to_frames(
    frame: Frame,
    alpha: torch.Tensor,
    aatype: torch.Tensor,
    default_frames: torch.Tensor,
):
    default_frame = Frame.from_tensor_4x4(default_frames[aatype, ...])

    bb_rot = alpha.new_zeros((*((1,) * len(alpha.shape[:-1])), 2))
    bb_rot[..., 1] = 1

    alpha = torch.cat([bb_rot.expand(*alpha.shape[:-2], -1, -1), alpha], dim=-2)

    all_rots = alpha.new_zeros(default_frame.get_rots().rot_mat.shape)
    all_rots[..., 0, 0] = 1
    all_rots[..., 1, 1] = alpha[..., 1]
    all_rots[..., 1, 2] = -alpha[..., 0]
    all_rots[..., 2, 1:] = alpha

    all_rots = Frame(Rotation(mat=all_rots), None)

    all_frames = default_frame.compose(all_rots)

    chi2_frame_to_frame = all_frames[..., 5]
    chi3_frame_to_frame = all_frames[..., 6]
    chi4_frame_to_frame = all_frames[..., 7]

    chi1_frame_to_bb = all_frames[..., 4]
    chi2_frame_to_bb = chi1_frame_to_bb.compose(chi2_frame_to_frame)
    chi3_frame_to_bb = chi2_frame_to_bb.compose(chi3_frame_to_frame)
    chi4_frame_to_bb = chi3_frame_to_bb.compose(chi4_frame_to_frame)

    all_frames_to_bb = Frame.cat(
        [
            all_frames[..., :5],
            chi2_frame_to_bb.unsqueeze(-1),
            chi3_frame_to_bb.unsqueeze(-1),
            chi4_frame_to_bb.unsqueeze(-1),
        ],
        dim=-1,
    )

    all_frames_to_global = frame[..., None].compose(all_frames_to_bb)

    return all_frames_to_global


def frames_and_literature_positions_to_atom14_pos(
    frame: Frame,
    aatype: torch.Tensor,
    default_frames,
    group_idx,
    atom_mask,
    lit_positions,
):
    group_mask = group_idx[aatype, ...]
    group_mask = one_hot(
        group_mask,
        num_classes=default_frames.shape[-3],
    )

    t_atoms_to_global = frame[..., None, :] * group_mask
    t_atoms_to_global = t_atoms_to_global.map_tensor_fn(lambda x: torch.sum(x, dim=-1))

    atom_mask = atom_mask[aatype, ...].unsqueeze(-1)

    lit_positions = lit_positions[aatype, ...]
    pred_positions = t_atoms_to_global.apply(lit_positions)
    pred_positions = pred_positions * atom_mask

    return pred_positions


class SideChainAngleResnetIteration(nn.Module):
    def __init__(self, d_hid):
        super(SideChainAngleResnetIteration, self).__init__()

        self.d_hid = d_hid

        self.linear_1 = Linear(self.d_hid, self.d_hid, init="relu")
        self.act = nn.GELU()
        self.linear_2 = Linear(self.d_hid, self.d_hid, init="final")

    def forward(self, s: torch.Tensor) -> torch.Tensor:

        x = self.act(s)
        x = self.linear_1(x)
        x = self.act(x)
        x = self.linear_2(x)

        return residual(s, x, self.training)


class SidechainAngleResnet(nn.Module):
    def __init__(self, d_in, d_hid, num_blocks, num_angles):
        super(SidechainAngleResnet, self).__init__()

        self.linear_in = Linear(d_in, d_hid)
        self.act = nn.GELU()
        self.linear_initial = Linear(d_in, d_hid)

        self.layers = SimpleModuleList()
        for _ in range(num_blocks):
            self.layers.append(SideChainAngleResnetIteration(d_hid=d_hid))

        self.linear_out = Linear(d_hid, num_angles * 2)

    def forward(
        self, s: torch.Tensor, initial_s: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:

        initial_s = self.linear_initial(self.act(initial_s))
        s = self.linear_in(self.act(s))

        s = s + initial_s

        for layer in self.layers:
            s = layer(s)

        s = self.linear_out(self.act(s))

        s = s.view(s.shape[:-1] + (-1, 2))

        unnormalized_s = s
        norm_denom = torch.sqrt(
            torch.clamp(
                torch.sum(s.float() ** 2, dim=-1, keepdim=True),
                min=1e-12,
            )
        )
        s = s.float() / norm_denom

        return unnormalized_s, s.type(unnormalized_s.dtype)


class InvariantPointAttention(nn.Module):
    def __init__(
        self,
        d_single: int,
        d_pair: int,
        d_hid: int,
        num_heads: int,
        num_qk_points: int,
        num_v_points: int,
        separate_kv: bool = False,
        bias: bool = True,
        eps: float = 1e-8,
    ):
        super(InvariantPointAttention, self).__init__()

        self.d_hid = d_hid
        self.num_heads = num_heads
        self.num_qk_points = num_qk_points
        self.num_v_points = num_v_points
        self.eps = eps

        hc = self.d_hid * self.num_heads
        self.linear_q = Linear(d_single, hc, bias=bias)
        self.separate_kv = separate_kv
        if self.separate_kv:
            self.linear_k = Linear(d_single, hc, bias=bias)
            self.linear_v = Linear(d_single, hc, bias=bias)
        else:
            self.linear_kv = Linear(d_single, 2 * hc, bias=bias)

        hpq = self.num_heads * self.num_qk_points * 3
        self.linear_q_points = Linear(d_single, hpq)
        hpk = self.num_heads * self.num_qk_points * 3
        hpv = self.num_heads * self.num_v_points * 3
        if self.separate_kv:
            self.linear_k_points = Linear(d_single, hpk)
            self.linear_v_points = Linear(d_single, hpv)
        else:
            hpkv = hpk + hpv
            self.linear_kv_points = Linear(d_single, hpkv)

        self.linear_b = Linear(d_pair, self.num_heads)

        self.head_weights = nn.Parameter(torch.zeros((num_heads)))
        ipa_point_weights_init_(self.head_weights)

        concat_out_dim = self.num_heads * (d_pair + self.d_hid + self.num_v_points * 4)
        self.linear_out = Linear(concat_out_dim, d_single, init="final")

        self.softplus = nn.Softplus()

    def forward(
        self,
        s: torch.Tensor,
        z: torch.Tensor,
        f: Frame,
        square_mask: torch.Tensor,
    ) -> torch.Tensor:
        q = self.linear_q(s) # [b,r,c]

        q = q.view(q.shape[:-1] + (self.num_heads, -1)) #[b,r,h,c]

        if self.separate_kv:
            k = self.linear_k(s)
            v = self.linear_v(s)
            k = k.view(k.shape[:-1] + (self.num_heads, -1))
            v = v.view(v.shape[:-1] + (self.num_heads, -1))
        else:
            kv = self.linear_kv(s)
            kv = kv.view(kv.shape[:-1] + (self.num_heads, -1))
            k, v = torch.split(kv, self.d_hid, dim=-1)

        q_pts = self.linear_q_points(s)

        def process_points(pts, no_points):
            shape = pts.shape[:-1] + (pts.shape[-1] // 3, 3)
            if self.separate_kv:
                # alphafold-multimer uses different layout
                pts = pts.view(pts.shape[:-1] + (self.num_heads, no_points * 3))
            pts = torch.split(pts, pts.shape[-1] // 3, dim=-1)
            pts = torch.stack(pts, dim=-1).view(*shape)
            pts = f[..., None].apply(pts)

            pts = pts.view(pts.shape[:-2] + (self.num_heads, no_points, 3))
            return pts

        q_pts = process_points(q_pts, self.num_qk_points)

        if self.separate_kv:
            k_pts = self.linear_k_points(s)
            v_pts = self.linear_v_points(s)
            k_pts = process_points(k_pts, self.num_qk_points)
            v_pts = process_points(v_pts, self.num_v_points)
        else:
            kv_pts = self.linear_kv_points(s)

            kv_pts = torch.split(kv_pts, kv_pts.shape[-1] // 3, dim=-1)
            kv_pts = torch.stack(kv_pts, dim=-1)
            kv_pts = f[..., None].apply(kv_pts)

            kv_pts = kv_pts.view(kv_pts.shape[:-2] + (self.num_heads, -1, 3))

            k_pts, v_pts = torch.split(
                kv_pts, [self.num_qk_points, self.num_v_points], dim=-2
            )

        bias = self.linear_b(z)#[b, h, q, k]

        #[b, h, q, k]
        attn = torch.matmul(
            permute_final_dims(q, (1, 0, 2)), #[b, h, r, c]
            permute_final_dims(k, (1, 2, 0)),#[b, h, c, r]
        )

        # q = q[:,:200]
        # bias = bias[:,:200]
        # bias = bias[:,:,:200]

        if self.training:
            attn = attn * math.sqrt(1.0 / (3 * self.d_hid))
            attn = attn + (math.sqrt(1.0 / 3) * permute_final_dims(bias, (2, 0, 1)))
            pt_att = q_pts.unsqueeze(-4) - k_pts.unsqueeze(-5)
            pt_att = pt_att.float() ** 2
        else:
            attn *= math.sqrt(1.0 / (3 * self.d_hid))
            attn += (math.sqrt(1.0 / 3) * permute_final_dims(bias, (2, 0, 1)))
            pt_att = q_pts.unsqueeze(-4) - k_pts.unsqueeze(-5)
            pt_att *= pt_att

        pt_att = pt_att.sum(dim=-1)
        head_weights = self.softplus(self.head_weights).view(
            *((1,) * len(pt_att.shape[:-2]) + (-1, 1))
        )
        head_weights = head_weights * math.sqrt(
            1.0 / (3 * (self.num_qk_points * 9.0 / 2))
        )
        pt_att *= head_weights * (-0.5)

        pt_att = torch.sum(pt_att, dim=-1)

        pt_att = permute_final_dims(pt_att, (2, 0, 1))
        attn += square_mask
        attn = softmax_dropout(attn, 0, self.training, bias=pt_att.type(attn.dtype))
        del pt_att, q_pts, k_pts, bias
        o = torch.matmul(attn, v.transpose(-2, -3)).transpose(-2, -3)
        o = o.contiguous().view(*o.shape[:-2], -1)
        del q, k, v
        o_pts = torch.sum(
            (
                attn[..., None, :, :, None]
                * permute_final_dims(v_pts, (1, 3, 0, 2))[..., None, :, :]
            ),
            dim=-2,
        )

        o_pts = permute_final_dims(o_pts, (2, 0, 3, 1))
        o_pts = f[..., None, None].invert_apply(o_pts)
        if self.training:
            o_pts_norm = torch.sqrt(torch.sum(o_pts.float() ** 2, dim=-1) + self.eps).type(
                o_pts.dtype
            )
        else:
            o_pts_norm = torch.sqrt(torch.sum(o_pts ** 2, dim=-1) + self.eps).type(
                o_pts.dtype
            )

        o_pts_norm = o_pts_norm.view(*o_pts_norm.shape[:-2], -1)

        o_pts = o_pts.view(*o_pts.shape[:-3], -1, 3)

        o_pair = torch.matmul(attn.transpose(-2, -3), z)

        o_pair = o_pair.view(*o_pair.shape[:-2], -1)

        s = self.linear_out(
            torch.cat((o, *torch.unbind(o_pts, dim=-1), o_pts_norm, o_pair), dim=-1)
        )

        return s
    
class ResNN(nn.Module):
    def __init__(
        self,
        d_single: int,
        d_pair: int,
        d_hid: int,
        num_heads: int,
        num_qk_points: int,
        num_v_points: int,
        separate_kv: bool = False,
        bias: bool = True,
        eps: float = 1e-8,
    ):
        super(ResNN, self).__init__()

        self.d_hid = d_hid
        self.num_heads = num_heads
        self.num_qk_points = num_qk_points
        self.num_v_points = num_v_points
        self.eps = eps

        hc = self.d_hid * self.num_heads
        self.linear_s = Linear(d_single, hc, bias=bias)
        self.linear_angle = Linear(d_pair, hc, bias=bias)
        self.linear_dist = Linear(d_pair, hc, bias=bias)
        self.linear_rel_pos = Linear(d_pair, hc, bias=bias)
        self.linear_merge_ij = Linear(hc*2, hc, bias=bias)
        self.ffn_merge = Transition(hc,1)
        self.ffn_dist = Transition(hc,1)

        self.linear_attn = Linear(hc, num_heads, init='final')
        self.linear_value = Linear(hc, hc)


        self.separate_kv = separate_kv

        hpq = self.num_heads * self.num_qk_points * 3
        hpk = self.num_heads * self.num_qk_points * 3
        hpv = self.num_heads * self.num_v_points * 3

        # self.head_weights = nn.Parameter(torch.zeros((num_heads)))
        # ipa_point_weights_init_(self.head_weights)

        # concat_out_dim = self.num_heads * (d_pair + self.d_hid + self.num_v_points * 4)
        self.linear_out = Linear(hc, d_single, init="final")


    def forward(
        self,
        s: torch.Tensor,
        pair_rep_angle: torch.Tensor,
        pair_rep_dist: torch.Tensor,
        relpos_emb: torch.Tensor,
        f: Frame,
        square_mask: torch.Tensor,
    ) -> torch.Tensor:
        s = self.linear_s(s) # [b,r,c]
        pair_rep_angle = self.linear_angle(pair_rep_angle)
        pair_rep_dist = self.linear_dist(pair_rep_dist)
        relpos_emb = self.linear_rel_pos(relpos_emb)

        pair_rep_angle = s[:,:,None,:] * pair_rep_angle # TODO weian: try +
        pair_rep_angle_inverse = rearrange(pair_rep_angle, 'b q k c -> b k q c')
        pair_rep_angle_merged = torch.cat([pair_rep_angle,pair_rep_angle_inverse],dim=-1)
        pair_rep_angle_merged = self.linear_merge_ij(pair_rep_angle_merged)
        pair_rep_angle_merged = residual(pair_rep_angle_merged,self.ffn_merge(pair_rep_angle_merged),training=self.training)
        pair_rep_angle_merged_dist = pair_rep_angle_merged + pair_rep_dist + relpos_emb
        pair_rep_angle_merged = residual(pair_rep_angle_merged,self.ffn_dist(pair_rep_angle_merged_dist),training=self.training)

        attn = self.linear_attn(pair_rep_angle_merged)
        value = self.linear_value(pair_rep_angle_merged)

        attn = rearrange(attn, 'b q k h -> b h q k')
        value = rearrange(value, 'b q k (h c) -> b h c q k', h = self.num_heads)


        # bias = self.linear_b(z)#[b, h, q, k]

        #[b, h, q, k]
        # attn = torch.matmul(
        #     permute_final_dims(q, (1, 0, 2)), #[b, h, r, c]
        #     permute_final_dims(k, (1, 2, 0)),#[b, h, c, r]
        # )

        # q = q[:,:200]
        # bias = bias[:,:200]
        # bias = bias[:,:,:200]


        attn += square_mask
        attn = softmax_dropout(attn, 0, self.training)
        o = torch.einsum('bhqv,bhcqv->bhcq',attn,value)
        o = rearrange(o,'b h c q -> b q (h c)')



        s = self.linear_out(o)

        return s


class BackboneUpdate(nn.Module):
    def __init__(self, d_single):
        super(BackboneUpdate, self).__init__()
        self.linear = Linear(d_single, 6, init="final")

    def forward(self, s: torch.Tensor):
        return self.linear(s)


class StructureModuleTransitionLayer(nn.Module):
    def __init__(self, c):
        super(StructureModuleTransitionLayer, self).__init__()

        self.linear_1 = Linear(c, c, init="relu")
        self.linear_2 = Linear(c, c, init="relu")
        self.act = nn.GELU()
        self.linear_3 = Linear(c, c, init="final")

    def forward(self, s):
        s_old = s
        s = self.linear_1(s)
        s = self.act(s)
        s = self.linear_2(s)
        s = self.act(s)
        s = self.linear_3(s)

        s = residual(s_old, s, self.training)

        return s


class StructureModuleTransition(nn.Module):
    def __init__(self, c, num_layers, dropout_rate):
        super(StructureModuleTransition, self).__init__()

        self.num_layers = num_layers
        self.dropout_rate = dropout_rate

        self.layers = SimpleModuleList()
        for _ in range(self.num_layers):
            self.layers.append(StructureModuleTransitionLayer(c))

        self.dropout = nn.Dropout(self.dropout_rate)
        self.layer_norm = LayerNorm(c)

    def forward(self, s):
        for layer in self.layers:
            s = layer(s)

        s = self.dropout(s)
        s = self.layer_norm(s)

        return s


class AFInverseGraghInitModule(nn.Module):
    def __init__(
        self,
        d_single,
        d_pair,
        d_ipa,
        d_angle,
        num_heads_ipa,
        num_qk_points,
        num_v_points,
        dropout_rate,
        num_blocks,
        no_transition_layers,
        num_resnet_blocks,
        num_angles,
        trans_scale_factor,
        separate_kv,
        ipa_bias,
        epsilon,
        inf,
        # use_structure_module_output,
        **kwargs,
    ):
        super(AFInverseGraghInitModule, self).__init__()

        self.num_blocks = num_blocks
        self.trans_scale_factor = trans_scale_factor
        self.default_frames = None
        self.group_idx = None
        self.atom_mask = None
        self.lit_positions = None
        self.inf = inf
        # self.use_structure_module_output = use_structure_module_output

        self.layer_norm_s = LayerNorm(d_single)
        # self.layer_norm_z = LayerNorm(d_pair)
        self.layer_norm_pair_angle = LayerNorm(d_pair)
        self.layer_norm_pair_dist = LayerNorm(d_pair)
        self.layer_norm_rel_pos = LayerNorm(d_pair)

        self.linear_in = Linear(d_single, d_single)

        # self.ipa = InvariantPointAttention(
        #     d_single,
        #     d_pair,
        #     d_ipa,
        #     num_heads_ipa,
        #     num_qk_points,
        #     num_v_points,
        #     separate_kv=separate_kv,
        #     bias=ipa_bias,
        #     eps=epsilon,
        # )
        self.resNN = ResNN(
            d_single,
            d_pair,
            d_ipa,
            num_heads_ipa,
            num_qk_points,
            num_v_points,
            separate_kv=separate_kv,
            bias=ipa_bias,
            eps=epsilon,
        )

        self.ipa_dropout = nn.Dropout(dropout_rate)
        self.layer_norm_ipa = LayerNorm(d_single)

        self.transition = StructureModuleTransition(
            d_single,
            no_transition_layers,
            dropout_rate,
        )
        assert len(rc.restype_1to3) == 20
        
        # self.predictor = Linear(d_single, 20, init = "final")
        self.predictor = InverseClsHead(num_res_type = 20, d_single = d_single, d_hid = d_single)



    def forward(
        self,
        s,
        z,
        gt_frame,
        mask=None,
    ):
        # s = s.type(self.linear_in.weight.dtype)
        # z = z.type(self.linear_in.weight.dtype)
        if mask is None:
            mask = s.new_ones(s.shape[:-1])

        # generate square mask
        square_mask = mask.unsqueeze(-1) * mask.unsqueeze(-2)
        square_mask = gen_attn_mask(square_mask, -self.inf).unsqueeze(-3)
        s = self.layer_norm_s(s)
        # z = self.layer_norm_z(z)
        # initial_s = s
        pair_rep_angle = self.layer_norm_pair_angle(z['pair_rep_angle'])
        pair_rep_dist = self.layer_norm_pair_dist(z['pair_rep_dist'])
        relpos_emb = self.layer_norm_rel_pos(z['relpos_emb'])
        s = self.linear_in(s)

        backb_to_global = gt_frame
        outputs = []
        for i in range(self.num_blocks):
            s_updated = self.resNN(s, pair_rep_angle, pair_rep_dist, relpos_emb, backb_to_global, square_mask) # TODO weian: updateing z
            s = residual(s, s_updated, self.training)
            s = self.ipa_dropout(s)
            s = self.layer_norm_ipa(s)
            s = self.transition(s) #TODO weian: remove unknow res type in loss

            res_score = self.predictor(s)
            preds = {
                "res_score": res_score,
            }

            outputs.append(preds)

        outputs = dict_multimap(torch.stack, outputs)

        return outputs

    def _init_residue_constants(self, float_dtype, device):
        if self.default_frames is None:
            self.default_frames = torch.tensor(
                restype_rigid_group_default_frame,
                dtype=float_dtype,
                device=device,
                requires_grad=False,
            )
        if self.group_idx is None:
            self.group_idx = torch.tensor(
                restype_atom14_to_rigid_group,
                device=device,
                requires_grad=False,
            )
        if self.atom_mask is None:
            self.atom_mask = torch.tensor(
                restype_atom14_mask,
                dtype=float_dtype,
                device=device,
                requires_grad=False,
            )
        if self.lit_positions is None:
            self.lit_positions = torch.tensor(
                restype_atom14_rigid_group_positions,
                dtype=float_dtype,
                device=device,
                requires_grad=False,
            )

    def torsion_angles_to_frames(self, frame, alpha, aatype):
        self._init_residue_constants(alpha.dtype, alpha.device)
        return torsion_angles_to_frames(frame, alpha, aatype, self.default_frames)

    def frames_and_literature_positions_to_atom14_pos(self, frame, aatype):
        self._init_residue_constants(frame.get_rots().dtype, frame.get_rots().device)
        return frames_and_literature_positions_to_atom14_pos(
            frame,
            aatype,
            self.default_frames,
            self.group_idx,
            self.atom_mask,
            self.lit_positions,
        )
