from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


torch.manual_seed(7)

FEATURE_COUNT = 6
BATCH_SIZE = 32
EPOCHS = 40
CHECKPOINT = Path("anomaly_detector.pt")


class Autoencoder(nn.Module):
    def __init__(self, feature_count: int) -> None:
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(feature_count, 4),
            nn.ReLU(),
            nn.Linear(4, 2),
        )
        self.decoder = nn.Sequential(
            nn.Linear(2, 4),
            nn.ReLU(),
            nn.Linear(4, feature_count),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.decoder(self.encoder(inputs))


def score_samples(model: nn.Module, samples: torch.Tensor) -> torch.Tensor:
    model.eval()
    with torch.inference_mode():
        reconstructed = model(samples)
        return torch.mean((samples - reconstructed) ** 2, dim=1)


def main() -> None:
    normal_samples = 0.20 * torch.randn(512, FEATURE_COUNT)
    loader = DataLoader(
        TensorDataset(normal_samples),
        batch_size=BATCH_SIZE,
        shuffle=True,
    )

    model = Autoencoder(FEATURE_COUNT)
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    for epoch in range(1, EPOCHS + 1):
        model.train()
        running_loss = 0.0

        for (batch,) in loader:
            optimizer.zero_grad()
            reconstructed = model(batch)
            loss = loss_fn(reconstructed, batch)
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * batch.size(0)

        if epoch in (1, 10, 20, 30, 40):
            print(f"epoch={epoch:02d} train_loss={running_loss / len(normal_samples):.6f}")

    train_scores = score_samples(model, normal_samples)
    threshold = train_scores.mean() + 3 * train_scores.std()

    normal_probe = 0.20 * torch.randn(16, FEATURE_COUNT)
    outlier_probe = normal_probe + 3.0
    normal_score_mean = score_samples(model, normal_probe).mean()
    outlier_score_mean = score_samples(model, outlier_probe).mean()

    torch.save(
        {
            "feature_count": FEATURE_COUNT,
            "threshold": threshold,
            "model_state": model.state_dict(),
        },
        CHECKPOINT,
    )

    print(f"threshold={threshold:.6f}")
    print(f"normal_score_mean={normal_score_mean:.6f}")
    print(f"outlier_score_mean={outlier_score_mean:.6f}")
    decision = bool(outlier_score_mean > threshold and normal_score_mean < threshold)
    print(f"anomaly_decision={decision}")
    print(f"checkpoint={CHECKPOINT}")


if __name__ == "__main__":
    main()