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_ts import FelaTsConfig
from .modeling import FELA_TS, FNO1D


def _fno1d_forward(self, x):
    w = torch.view_as_complex(self.w)
    P = x.shape[1]
    xf = torch.fft.rfft(x, dim=1)
    mm = min(self.m, xf.shape[1])
    o = torch.zeros_like(xf)
    o[:, :mm] = torch.einsum("bpd,pde->bpe", xf[:, :mm], w[:mm])
    return torch.fft.irfft(o, n=P, dim=1)


def _realify(model):
    for m in model.modules():
        if isinstance(m, FNO1D):
            m.w = nn.Parameter(torch.view_as_real(m.w.detach()).contiguous())
            m.forward = types.MethodType(_fno1d_forward, m)


class FelaTsModel(PreTrainedModel):
    config_class = FelaTsConfig
    base_model_prefix = "model"
    main_input_name = "x"

    def __init__(self, config):
        super().__init__(config)
        self.model = FELA_TS(
            config.C,
            config.L,
            config.H,
            config.patch,
            config.stride,
            config.D,
            config.modes,
            config.nblk,
        )
        _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)