from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F


class InferenceGraphSAGE(nn.Module):
    def __init__(self, in_dim: int = 768, hidden_dim: int = 256, num_layers: int = 2, dropout: float = 0.1):
        super().__init__()
        self.dropout = float(dropout)

        self.linears = nn.ModuleList()
        cur_in = in_dim
        for _ in range(num_layers):
            self.linears.append(nn.Linear(cur_in * 2, hidden_dim))
            cur_in = hidden_dim

        self.class_wheeze = nn.Linear(cur_in, 1)
        self.class_crackle = nn.Linear(cur_in, 1)
        self.act = nn.ReLU()

    def forward(self, x: torch.Tensor, adj: torch.Tensor):
        for lin in self.linears:
            deg = adj.sum(dim=-1, keepdim=True).clamp(min=1.0)
            neighbor_mean = torch.matmul(adj, x) / deg
            h = torch.cat([x, neighbor_mean], dim=-1)
            h = self.act(lin(h))
            h = F.dropout(h, p=self.dropout, training=self.training)
            x = h

        wheeze = torch.sigmoid(self.class_wheeze(x).squeeze(-1))
        crackle = torch.sigmoid(self.class_crackle(x).squeeze(-1))
        return wheeze, crackle


def load_checkpoint_flexible(model: nn.Module, path: str, map_location: Optional[str] = "cpu") -> nn.Module:
    ck = torch.load(path, map_location=map_location)
    if isinstance(ck, dict) and ("state_dict" in ck or "model_state_dict" in ck):
        state = ck.get("state_dict", ck.get("model_state_dict"))
    else:
        state = ck

    try:
        model.load_state_dict(state, strict=False)
        return model
    except Exception:
        ms = model.state_dict()
        incoming = {}
        for k, v in state.items():
            k2 = k[7:] if k.startswith("module.") else k
            incoming[k2] = v

        resolved = {}
        for k, v in incoming.items():
            if k in ms and ms[k].shape == v.shape:
                resolved[k] = v

        for k_in, v in incoming.items():
            if k_in in resolved:
                continue
            for k_model, v_model in ms.items():
                if k_model in resolved:
                    continue
                if v_model.shape == v.shape and k_model.endswith(k_in.split(".")[-1]):
                    resolved[k_model] = v
                    break

        ms.update(resolved)
        model.load_state_dict(ms)
        return model