import os
import sys
import types

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

from .configuration_pde import FelaPdeConfig
from .modeling import FNO2d, SpectralConv2d


def _spectral_forward(self, x):
    w1 = torch.view_as_complex(self.w1)
    w2 = torch.view_as_complex(self.w2)
    B, C, Hh, Ww = x.shape
    xf = torch.fft.rfft2(x)
    o = torch.zeros(
        B, w1.shape[1], Hh, Ww // 2 + 1, dtype=torch.cfloat, device=x.device
    )
    o[:, :, : self.m1, : self.m2] = torch.einsum(
        "bixy,ioxy->boxy", xf[:, :, : self.m1, : self.m2], w1
    )
    o[:, :, -self.m1 :, : self.m2] = torch.einsum(
        "bixy,ioxy->boxy", xf[:, :, -self.m1 :, : self.m2], w2
    )
    return torch.fft.irfft2(o, s=(Hh, Ww))


def _realify(model):
    for m in model.modules():
        if isinstance(m, SpectralConv2d):
            m.w1 = nn.Parameter(torch.view_as_real(m.w1.detach()).contiguous())
            m.w2 = nn.Parameter(torch.view_as_real(m.w2.detach()).contiguous())
            m.forward = types.MethodType(_spectral_forward, m)


class FelaPdeModel(PreTrainedModel):
    config_class = FelaPdeConfig
    base_model_prefix = "model"
    main_input_name = "x"

    def __init__(self, config):
        super().__init__(config)
        self.model = FNO2d(
            in_ch=config.in_ch,
            modes=config.modes,
            width=config.width,
            L=config.layers,
            proj_hidden=config.proj_hidden,
        )
        _realify(self.model)
        self.post_init()

    def forward(self, x=None, input_values=None, **kwargs):
        if x is None:
            x = input_values
        out = self.model(x)
        return CausalLMOutput(logits=out)