"""Per-pixel spectral time-series CNNs: original mean/max and attention/max variants."""
import torch
from torch import nn

def valid_from_index_positions(pos):
    """The unmodified dataset uses 0,1,...,T-1 and pads positions with zero.

Temporal subsampling may drop acquisition zero; its first retained position is
then positive. Only an initial zero is a valid acquisition-zero token.
This cannot be reused for arbitrary calendar positions without a length mask.
"""
    first_is_zero = pos[:, :1].eq(0)
    first = torch.zeros_like(pos, dtype=torch.bool)
    first[:, :1] = first_is_zero
    valid = pos.gt(0) | first
    if not valid.any(dim=1).all():
        raise ValueError('Each time series must contain an observation')
    return valid


class TemporalCNN(nn.Module):
    def __init__(self, classes=18):
        super().__init__()
        self.input=nn.Linear(10,64)
        self.norms=nn.ModuleList([nn.LayerNorm(64) for _ in range(4)])
        self.convs=nn.ModuleList([nn.Conv1d(64,64,3,padding=d,dilation=d,groups=64) for d in [1,2,4,8]])
        self.mix=nn.ModuleList([nn.Linear(64,64) for _ in range(4)])
        self.head=nn.Sequential(nn.Linear(128,64),nn.GELU(),nn.Linear(64,classes))

    def forward(self,x,pos):
        b,t,c,h,w=x.shape
        valid=valid_from_index_positions(pos)[:,None,None].expand(b,h,w,t).reshape(-1,t,1)
        flat=x.permute(0,3,4,1,2).reshape(-1,t,c)
        outputs=[]
        for start in range(0,len(flat),2048):
            mask=valid[start:start+2048];z=self.input(flat[start:start+2048])*mask
            for norm,conv,mix in zip(self.norms,self.convs,self.mix):
                z=(z+mix(torch.nn.functional.gelu(conv(norm(z).transpose(1,2)).transpose(1,2))))*mask
            mean=z.sum(1)/mask.sum(1).clamp_min(1)
            maximum=z.masked_fill(~mask,torch.finfo(z.dtype).min).amax(1)
            outputs.append(self.head(torch.cat([mean,maximum],-1)))
        return torch.cat(outputs).reshape(b,h,w,-1).permute(0,3,1,2)


class AttentionCNN(TemporalCNN):
    def __init__(self,classes=18):
        super().__init__(classes);self.attention=nn.Linear(64,1)
    def forward(self,x,pos):
        b,t,c,h,w=x.shape;valid=valid_from_index_positions(pos)[:,None,None].expand(b,h,w,t).reshape(-1,t,1);flat=x.permute(0,3,4,1,2).reshape(-1,t,c);out=[]
        for start in range(0,len(flat),2048):
            mask=valid[start:start+2048];z=self.input(flat[start:start+2048])*mask
            for norm,conv,mix in zip(self.norms,self.convs,self.mix):z=(z+mix(torch.nn.functional.gelu(conv(norm(z).transpose(1,2)).transpose(1,2))))*mask
            score=self.attention(z).masked_fill(~mask,torch.finfo(z.dtype).min)
            pooled=(score.softmax(1)*z).sum(1);maximum=z.masked_fill(~mask,torch.finfo(z.dtype).min).amax(1)
            out.append(self.head(torch.cat([pooled,maximum],-1)))
        return torch.cat(out).reshape(b,h,w,-1).permute(0,3,1,2)


def make_model(kind,classes=18):
    return {"cnn_control":TemporalCNN,"attention":AttentionCNN}[kind](classes)