"""
Standalone architecture definition for the Cluster 1 Similitude Regressor.

This is a self-contained copy of the model class used by the
Scientific AI Cluster Orchestration Framework's Cluster 1 pipeline
(https://huggingface.co/spaces/dave1368/cluster-01-dimensional-analysis) --
included here so the checkpoint can be loaded independently of that app.

Usage:
    from huggingface_hub import hf_hub_download
    from modeling import SimilitudeRegressorPINN
    import torch

    ckpt_path = hf_hub_download("dave1368/cluster-01-similitude-regressor", "similitude_regressor.pt")
    checkpoint = torch.load(ckpt_path, map_location="cpu")

    model = SimilitudeRegressorPINN()
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    # Inputs are RAW [Reynolds, Froude, Mach] -- normalization happens inside forward().
    pi_groups = torch.tensor([[10000.0, 0.5, 0.2]])  # Re=10,000, Fr=0.5, Mach=0.2
    cd, cp = model(pi_groups)[0].tolist()
    print(f"Cd={cd:.4f}  Cp={cp:.4f}")
"""
import torch
import torch.nn as nn

# Re spans orders of magnitude (10 to 1e6); log-scaling it keeps the calling
# convention "pass raw [Re, Fr, Mach]" simple for callers.
RE_LOG_MIN, RE_LOG_MAX = 1.0, 6.0  # log10(10) .. log10(1e6)


class SimilitudeRegressorPINN(nn.Module):
    """
    Neural regressor mapping dimensionless input groups (Re, Fr, Mach)
    to scaled performance coefficients (Cd, Cp).

    Trained against Morrison (2013)'s sphere drag correlation Cd(Re) and the
    Prandtl-Glauert compressibility-corrected stagnation Cp(Mach). See this
    repo's README.md for the full training story, including the drag-crisis
    oversampling fix.
    """
    def __init__(self, input_dim=3, hidden_neurons=64):
        super().__init__()
        layers = [nn.Linear(input_dim, hidden_neurons), nn.Tanh()]
        for _ in range(3):
            layers.extend([nn.Linear(hidden_neurons, hidden_neurons), nn.Tanh()])

        # Output: [Drag Coefficient Cd, Pressure Coefficient Cp]
        layers.append(nn.Linear(hidden_neurons, 2))
        self.net = nn.Sequential(*layers)

    def _normalize(self, pi_groups: torch.Tensor) -> torch.Tensor:
        log_re = torch.log10(torch.clamp(pi_groups[:, 0], min=1.0))
        re_norm = 2.0 * (log_re - RE_LOG_MIN) / (RE_LOG_MAX - RE_LOG_MIN) - 1.0
        fr_norm = 2.0 * pi_groups[:, 1] - 1.0
        mach_norm = 2.0 * pi_groups[:, 2] - 1.0
        return torch.stack([re_norm, fr_norm, mach_norm], dim=1)

    def forward(self, pi_groups: torch.Tensor) -> torch.Tensor:
        return self.net(self._normalize(pi_groups))