import os
import torch

from torch import nn
from io import BytesIO
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad

def decrypt_model(configs, input_path):
    with open(input_path, "rb") as f:
        data = f.read()

    with open(
        os.path.join(configs["binary_path"], "decrypt.bin"), 
        "rb"
    ) as f:
        key = f.read()

    return BytesIO(
        unpad(
            AES.new(
                key, 
                AES.MODE_CBC, 
                data[:16]
            ).decrypt(data[16:]), 
            AES.block_size
        )
    ).read()

def calc_same_padding(kernel_size):
    pad = kernel_size // 2
    return (pad, pad - (kernel_size + 1) % 2)

def torch_interp(x, xp, fp):
    sort_idx = xp.argsort()
    xp = xp[sort_idx]
    fp = fp[sort_idx]

    right_idxs = torch.searchsorted(xp, x).clamp(max=len(xp) - 1)
    left_idxs = (right_idxs - 1).clamp(min=0)
    x_left = xp[left_idxs]
    y_left = fp[left_idxs]

    interp_vals = y_left + ((x - x_left) * (fp[right_idxs] - y_left) / (xp[right_idxs] - x_left))
    interp_vals[x < xp[0]] = fp[0]
    interp_vals[x > xp[-1]] = fp[-1]

    return interp_vals

def batch_interp_with_replacement_detach(uv, f0):
    result = f0.clone()

    for i in range(uv.shape[0]):
        interp_vals = torch_interp(
            torch.where(uv[i])[-1], 
            torch.where(~uv[i])[-1], 
            f0[i][~uv[i]]
        ).detach()

        result[i][uv[i]] = interp_vals

    return result

class DotDict(dict):
    def __getattr__(*args):
        val = dict.get(*args)
        return DotDict(val) if type(val) is dict else val

    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

class Swish(nn.Module):
    def forward(self, x):
        return x * x.sigmoid()

class Transpose(nn.Module):
    def __init__(self, dims):
        super().__init__()
        assert len(dims) == 2, "dims == 2"
        self.dims = dims

    def forward(self, x):
        return x.transpose(*self.dims)

class GLU(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        out, gate = x.chunk(2, dim=self.dim)
        return out * gate.sigmoid()