import os
import sys
from collections import defaultdict

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F

from mapping_network import MappingNetowrk, AdaptiveInstanceNorm, NoiseInjection
from helpers.imle_helpers import (
    get_1x1,
    get_3x3,
    draw_gaussian_diag_samples,
    gaussian_analytical_kl,
)
from rtm_core import RTMMappingNetwork


class Block(nn.Module):
    def __init__(self, in_width, middle_width, out_width, down_rate=None, residual=False, use_3x3=True, zero_last=False):
        super().__init__()
        self.down_rate = down_rate
        self.residual = residual
        self.c1 = get_1x1(in_width, middle_width)
        self.c2 = get_3x3(middle_width, middle_width) if use_3x3 else get_1x1(
            middle_width, middle_width)
        self.c3 = get_3x3(middle_width, middle_width) if use_3x3 else get_1x1(
            middle_width, middle_width)
        self.c4 = get_1x1(middle_width, out_width, zero_weights=zero_last)

    def forward(self, x):
        xhat = self.c1(F.gelu(x))
        xhat = self.c2(F.gelu(xhat))
        xhat = self.c3(F.gelu(xhat))
        xhat = self.c4(F.gelu(xhat))
        out = x + xhat if self.residual else xhat
        if self.down_rate is not None:
            out = F.avg_pool2d(
                out, kernel_size=self.down_rate, stride=self.down_rate)
        return out


def parse_layer_string(s):
    layers = []
    for ss in s.split(','):
        if 'x' in ss:
            res, num = ss.split('x')
            count = int(num)
            layers += [(int(res), None) for _ in range(count)]
        elif 'm' in ss:
            res, mixin = [int(a) for a in ss.split('m')]
            layers.append((res, mixin))
        elif 'd' in ss:
            res, down_rate = [int(a) for a in ss.split('d')]
            layers.append((res, down_rate))
        else:
            res = int(ss)
            layers.append((res, None))
    return layers


def pad_channels(t, width):
    d1, d2, d3, d4 = t.shape
    empty = torch.zeros(d1, width, d3, d4, device=t.device)
    empty[:, :d2, :, :] = t
    return empty


def get_width_settings(width, s):
    mapping = defaultdict(lambda: width)
    if s:
        s = s.split(',')
        for ss in s:
            k, v = ss.split(':')
            mapping[int(k)] = int(v)
    return mapping


class DecBlock(nn.Module):
    def __init__(self, H, res, mixin, n_blocks):
        super().__init__()
        self.base = res
        self.mixin = mixin
        self.H = H
        self.widths = get_width_settings(H.width, H.custom_width_str)
        width = self.widths[res]
        if res <= H.max_hierarchy:
            self.noise = NoiseInjection(width)
        self.adaIN = AdaptiveInstanceNorm(width, H.latent_dim)
        use_3x3 = res > 2
        cond_width = int(width * H.bottleneck_multiple)
        self.resnet = Block(width, cond_width, width,
                            residual=True, use_3x3=use_3x3)
        self.resnet.c4.weight.data *= np.sqrt(1 / n_blocks)

    def forward(self, x, w, spatial_noise):
        if self.mixin is not None:
            x = F.interpolate(x, scale_factor=self.base // self.mixin)
        if self.base <= self.H.max_hierarchy:
            x = self.noise(x, spatial_noise)
        x = self.adaIN(x, w)
        x = self.resnet(x)
        return x


class Decoder(nn.Module):
    def __init__(self, H):
        super().__init__()
        self.H = H
        self.use_rtm = getattr(H, "use_rtm", False)

        if self.use_rtm:
            self.mapping_network = RTMMappingNetwork(
                code_dim=H.latent_dim,
                num_tokens=getattr(H, "num_tokens", 1),
                H_cycles=getattr(H, "H_cycles", 1),
                L_cycles=getattr(H, "L_cycles", 1),
                H_layers=getattr(H, "H_layers", 2),
                L_layers=getattr(H, "L_layers", 2),
                hidden_size=getattr(H, "rtm_hidden_size", 256),
                expansion=getattr(H, "rtm_expansion", 4.0),
                refinement_steps=getattr(H, "refinement_steps", 1),
                with_grad=getattr(H, "rtm_with_grad", False),
                cycle_noise_std=getattr(H, "rtm_cycle_noise_std", 0.0),
            )
        else:
            self.mapping_network = MappingNetowrk(
                code_dim=H.latent_dim, n_mlp=H.n_mpl)

        resos = set()
        cond_width = int(H.width * H.bottleneck_multiple)
        dec_blocks = []
        self.widths = get_width_settings(H.width, H.custom_width_str)
        blocks = parse_layer_string(H.dec_blocks)
        for idx, (res, mixin) in enumerate(blocks):
            dec_blocks.append(DecBlock(H, res, mixin, n_blocks=len(blocks)))
            resos.add(res)
        self.resolutions = sorted(resos)
        self.dec_blocks = nn.ModuleList(dec_blocks)
        first_res = self.resolutions[0]
        self.constant = nn.Parameter(torch.randn(
            1, self.widths[first_res], first_res, first_res))
        self.resnet = get_1x1(H.width, H.image_channels)
        self.gain = nn.Parameter(torch.ones(1, H.image_channels, 1, 1))
        self.bias = nn.Parameter(torch.zeros(1, H.image_channels, 1, 1))

    def forward(self, latent_code, spatial_noise, input_is_w=False):
        if not input_is_w:
            ws = self.mapping_network(latent_code)
        else:
            ws = [latent_code]

        x = self.constant.repeat(latent_code.shape[0], 1, 1, 1)

        num_blocks = len(self.dec_blocks)
        if len(ws) == 1:
            w = ws[0]
            for idx, block in enumerate(self.dec_blocks):
                x = block(x, w, None)
        else:
            # Multi-scale w injection: reversed so refined (last cycle) goes
            # to early/low-res blocks, diverse (first cycle) to late/high-res.
            ws_rev = list(reversed(ws))
            num_ws = len(ws_rev)
            for idx, block in enumerate(self.dec_blocks):
                w_idx = min(int(idx * num_ws / num_blocks), num_ws - 1)
                x = block(x, ws_rev[w_idx], None)

        x = self.resnet(x)
        x = self.gain * x + self.bias
        return x


class IMLE(nn.Module):
    def __init__(self, H):
        super().__init__()
        self.dci_db = None
        self.decoder = Decoder(H)

    def forward(self, latents, spatial_noise=None, input_is_w=False):
        return self.decoder.forward(latents, spatial_noise, input_is_w)