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_grid import FelaGridConfig
from .modeling import FELA_Grid, 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.modes, 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 FelaGridModel(PreTrainedModel):
    config_class = FelaGridConfig
    base_model_prefix = "model"
    main_input_name = "x"

    def __init__(self, config):
        super().__init__(config)
        self.model = FELA_Grid(
            config.Fin,
            config.L,
            D=config.D,
            modes=config.modes,
            nblk=config.nblk,
            nq=config.nq,
            arch=config.arch,
        )
        self.model._non_persistent_buffers_set.add("qidx")
        _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)