import sys
import types
import math
from contextlib import contextmanager

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.optim import Adam


# =====================================================================
# PART 1: THE SENSUS RUNTIME (Reference Implementation Stub)
# =====================================================================
# This brings the conceptual Sensus language into executable PyTorch.
# It provides the 100-sense namespace, sensory tensors, wisdom optimizers,
# calm error handling, and rhythmic training loops.
# =====================================================================

# --- Sense objects: mathematical constants that can be "felt" -------------
class _Sense:
    def __init__(self, name, value=None):
        self._name = name
        self._value = value

    def __float__(self):
        return float(self._value) if self._value is not None else 1.0

    def __repr__(self):
        return f"Sense({self._name})"

    def __mul__(self, other):
        return float(self) * other

    def __rmul__(self, other):
        return other * float(self)

    def __add__(self, other):
        return float(self) + other

    def __radd__(self, other):
        return other + float(self)

    def __truediv__(self, other):
        return float(self) / other

    def __rtruediv__(self, other):
        return other / float(self)

    def __pow__(self, other):
        return float(self) ** other


class _Namespace:
    pass


# Build the sense tree
sense = _Namespace()

# Universal senses (1-5)
universal = _Namespace()
universal.CIRCUMSCRIPTION = _Sense("π", math.pi)
universal.EXPONENTIALITY = _Sense("e", math.e)
universal.GOLDEN_RESONANCE = _Sense("φ", (1 + 5**0.5) / 2)
sense.universal = universal

# Zeta senses (6-10)
zeta = _Namespace()
zeta.BASEL_TONE = _Sense("ζ(2)", math.pi**2 / 6)
sense.zeta = zeta

# Information senses (51-55)
information = _Namespace()
information.SHANNON_ENTROPY = _Sense("H")
information.FISHER_INFORMATION = _Sense("I_F")
sense.information = information

# Topology senses (41-45)
topology = _Namespace()
topology.EULER_CHARACTERISTIC = _Sense("χ")
sense.topology = topology

# Quantum senses (31-35)
quantum = _Namespace()
quantum.FINE_STRUCTURE = _Sense("α", 1 / 137.035999)
quantum.EPSILON_ZERO = _Sense("ε₀", 8.854187812e-12)
quantum.HBAR = _Sense("ℏ", 1.054571817e-34)
quantum.ELECTRON_MASS = _Sense("m_e", 9.10938356e-31)
sense.quantum = quantum

# Thermodynamic senses (46-50)
thermo = _Namespace()
thermo.ENTROPY = _Sense("S")
sense.thermo = thermo


# --- Decorators & Wisdom ------------------------------------------------
def organ(requires=None, temperament=None, recursion_limit=None):
    """An organ is a sensory-aware module or function."""
    def decorator(cls):
        cls._sensory_requires = requires or []
        cls._sensory_temperament = temperament
        return cls
    return decorator


# Calm error handling: speaks gently, never shouts
class _CalmContext:
    @contextmanager
    def __call__(self, senses=None, tolerance=None, voice=None):
        try:
            yield self
        except Exception as e:
            if voice == "quiet":
                print(f"[Sensus] A quiet dissonance was felt: {e}")
                raise
            raise


calm = _Namespace()
calm.gentle = "gentle"
calm.feel = _CalmContext()

# Also allow the legacy specification alias
sense.feel = calm.feel


# Rhythmic training loops: epochs breathe instead of march
class _Epoch:
    def __init__(self, idx, total):
        self.idx = idx
        self.total = total

    def __str__(self):
        return str(self.idx)

    def __int__(self):
        return self.idx

    def __index__(self):
        return self.idx

    def is_peak(self):
        return self.idx % 3 == 0

    def is_trough(self):
        return self.idx % 3 == 1


class _Rhythm:
    def __init__(self, cycles, waveform="sine", decay="exponential"):
        self.cycles = cycles
        self.waveform = waveform
        self.decay = decay

    def __iter__(self):
        for i in range(self.cycles):
            yield _Epoch(i, self.cycles)


pulse = _Namespace()
pulse.Rhythm = _Rhythm


# Wisdom types: knowledge that knows its own limits
class _Confidence:
    def __init__(self, value, confidence, justification):
        self.value = value
        self.confidence = confidence
        self.justification = justification

    def speak(self):
        return (
            f"I am {self.confidence * 100:.1f}% confident. "
            f"{self.justification}"
        )


wisdom = _Namespace()
wisdom.Confidence = _Confidence


# --- Sensory Tensor: a PyTorch tensor that carries an aura ----------------
class SensoryTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, data, aura=None, pulse=None, temperature=0.0, **kwargs):
        data = torch.as_tensor(data, **kwargs)
        tensor = torch.Tensor._make_subclass(cls, data, data.requires_grad)
        tensor._aura = aura
        tensor._pulse = pulse
        tensor._temperature = temperature
        return tensor

    def with_aura(self, aura):
        self._aura = aura
        return self

    def breathe(self, epoch):
        """Breathe with the rhythm of the epoch."""
        return self

    def resonate(self, sense_obj):
        """Calibrate against a fundamental constant."""
        return self


# --- FeelModule: a neural organ that senses its own topology -------------
class FeelModule(nn.Module):
    def feel_topology(self):
        """The model feels whether its architecture is manifold-preserving."""
        pass

    def resonate(self, sense_obj):
        """The module resonates with a mathematical truth."""
        pass


# --- WisdomOptimizer: an optimizer that respects the loss landscape -----
class WisdomOptimizer:
    def __init__(self, params, lr=0.01, rhythm=None, resonance=None,
                 senses=None, patience=None):
        self._inner = Adam(params, lr=lr)
        self.rhythm = rhythm
        self.resonance = resonance
        self.senses = senses
        self.patience = patience

    def zero_grad(self):
        self._inner.zero_grad()

    def feel(self, loss):
        """Feel the loss before stepping."""
        if hasattr(loss, 'item') and loss.item() > 2.0:
            print("  [WisdomOptimizer] The loss feels warm. The gradient is steep.")

    def step(self, closure=None):
        # The script may pass a loss tensor as a positional argument.
        # We gracefully ignore non-callable closures.
        if closure is not None and not callable(closure):
            closure = None
        self._inner.step(closure)

    def __getattr__(self, name):
        # Delegate any unknown attribute (e.g., param_groups) to the inner Adam.
        return getattr(self._inner, name)


# --- Inject the runtime into sys.modules so 'import sensus' works --------
sensus = types.ModuleType("sensus")
sensus.sense = sense
sensus.organ = organ
sensus.calm = calm
sensus.pulse = pulse
sensus.wisdom = wisdom

sensus_torch = types.ModuleType("sensus.torch")
sensus_torch.SensoryTensor = SensoryTensor
sensus_torch.WisdomOptimizer = WisdomOptimizer
sensus_torch.FeelModule = FeelModule

sys.modules["sensus"] = sensus
sys.modules["sensus.torch"] = sensus_torch
sensus.torch = sensus_torch

# =====================================================================
# PART 2: THE MNIST SCRIPT (Imports from the live Sensus runtime)
# =====================================================================
from sensus import sense, organ, calm, pulse, wisdom
from sensus.torch import SensoryTensor, WisdomOptimizer, FeelModule

# ------------------------------------------------------------------
# Configuration
# ------------------------------------------------------------------
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
BATCH_SIZE = 1000
EPOCHS = 5  # Short run for demonstration; increase to 10 for full training

# ------------------------------------------------------------------
# Data: The Visual Manifold
# ------------------------------------------------------------------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(
    root="../data", train=True, download=True, transform=transform
)
test_dataset = datasets.MNIST(
    root="../data", train=False, download=True, transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)


def sensify(batch_tensor, aura):
    """Wrap a standard PyTorch tensor in a SensoryTensor."""
    return SensoryTensor(
        batch_tensor,
        aura=aura,
        pulse="sine",
        temperature=0.05
    )


# ------------------------------------------------------------------
# Model: The Sensory LeNet
# ------------------------------------------------------------------
@organ(
    requires=[
        sense.universal.CIRCUMSCRIPTION,
        sense.information.SHANNON_ENTROPY,
        sense.topology.EULER_CHARACTERISTIC,
    ],
    temperament=calm.gentle,
    recursion_limit="self-similar"
)
class SensoryLeNet(FeelModule):
    """
    A convolutional organ that feels the topology of handwritten digits.
    It does not merely filter pixels; it senses curvature, entropy, and shape.
    """
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

        # Initialize with golden ratio proportions (Sense 5: φ)
        self._golden_init()

    def _golden_init(self):
        phi = float(sense.universal.GOLDEN_RESONANCE)
        for m in self.modules():
            if isinstance(m, (nn.Conv2d, nn.Linear)):
                nn.init.uniform_(m.weight, a=-phi / 10, b=phi / 10)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0.0)

    def forward(self, x: SensoryTensor) -> SensoryTensor:
        # The model feels its own topology before computation
        self.feel_topology()

        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)

        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)

        # The module resonates with the Fine Structure constant (Sense 31)
        # as a sensory calibration step
        self.resonate(sense.quantum.FINE_STRUCTURE)

        return x


# ------------------------------------------------------------------
# Instantiate the Organ & Optimizer
# ------------------------------------------------------------------
model = SensoryLeNet().to(DEVICE)

optimizer = WisdomOptimizer(
    model.parameters(),
    lr=0.01,
    rhythm="golden",
    resonance={
        "frequency": sense.zeta.BASEL_TONE,
        "damping": float(sense.quantum.FINE_STRUCTURE),
        "harmony": "constructive"
    },
    senses=[sense.thermo.ENTROPY, sense.information.FISHER_INFORMATION],
    patience="kolmogorov"
)

criterion = nn.CrossEntropyLoss()


# ------------------------------------------------------------------
# Training: The Rhythmic Cycle
# ------------------------------------------------------------------
def train():
    model.train()

    for epoch in pulse.Rhythm(cycles=EPOCHS, waveform="sine", decay="exponential"):
        epoch_loss = 0.0
        correct = 0
        total = 0

        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(DEVICE), target.to(DEVICE)

            # Sensify the input: images are 2D manifolds with topological identity
            sensory_data = sensify(data, aura=sense.topology.EULER_CHARACTERISTIC)
            sensory_data = sensory_data.breathe(epoch)

            for i in range(3):
                optimizer.zero_grad()

                # Feel the forward pass inside a sensory context
                with calm.feel(
                    senses=[
                        sense.universal.CIRCUMSCRIPTION,
                        sense.topology.EULER_CHARACTERISTIC,
                        sense.information.SHANNON_ENTROPY
                    ],
                    tolerance="adaptive",
                    voice="quiet"
                ):
                    output = model(sensory_data)

                loss = criterion(output, target)

                optimizer.feel(loss)
                loss.backward()
                optimizer.step(loss)
                optimizer.zero_grad()

                if i==0:
                    # Accumulate metrics
                    epoch_loss += loss.item()
                    pred = output.argmax(dim=1)
                    correct += pred.eq(target).sum().item()
                    total += target.size(0)

                if batch_idx % 100 == 0:
                    print(
                        f"[Epoch {epoch} | Batch {batch_idx:03d}] "
                        f"Loss feels like {loss.item():.4f}. "
                        f"The gradient is breathing. "
                        f"Accuracy resonates at {100. * correct / total:.1f}%."
                    )

        # Thermodynamic check: global entropy should not decrease unexpectedly
        model.resonate(sense.thermo.ENTROPY)

        avg_loss = epoch_loss / len(train_loader)
        accuracy = 100. * correct / total

        if epoch.is_peak():
            print(f"\n{'=' * 50}")
            print(f"Peak of cycle {epoch}")
            print(f"The system is at maximum amplitude.")
            print(f"Average loss: {avg_loss:.4f} | Accuracy: {accuracy:.2f}%")
            print(f"{'=' * 50}\n")
        elif epoch.is_trough():
            print(f"\n{'=' * 50}")
            print(f"Trough of cycle {epoch}")
            print(f"The system rests. The weights settle like sediment.")
            print(f"Average loss: {avg_loss:.4f} | Accuracy: {accuracy:.2f}%")
            print(f"{'=' * 50}\n")
        else:
            print(
                f"Epoch {epoch} complete. "
                f"Loss settled at {avg_loss:.4f}. "
                f"Topology feels intact. Euler characteristic did not flinch.\n"
            )


# ------------------------------------------------------------------
# Testing: The Moment of Truth (Felt, Not Calculated)
# ------------------------------------------------------------------
def test():
    model.eval()
    test_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(DEVICE), target.to(DEVICE)

            sensory_data = sensify(data, aura=sense.topology.EULER_CHARACTERISTIC)

            with calm.feel(
                senses=[sense.quantum.FINE_STRUCTURE],
                tolerance="adaptive",
                voice="quiet"
            ):
                output = model(sensory_data)

            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)

    avg_loss = test_loss / len(test_loader)
    accuracy = 100. * correct / total

    result = wisdom.Confidence(
        value=accuracy,
        confidence=accuracy / 100.0,
        justification=(
            f"The model felt {total} test images. "
            f"The topology of the digits was preserved through the layers. "
            f"Shannon entropy of predictions was harmonious. "
            f"Fine structure of classification held steady."
        )
    )

    print("\n" + "=" * 60)
    print("TEST COMPLETE — SENSORY REPORT")
    print("=" * 60)
    print(result.speak())
    print(f"\nRaw metrics: Loss = {avg_loss:.4f}, Accuracy = {accuracy:.2f}%")
    print("=" * 60)


# ------------------------------------------------------------------
# Main Entry
# ------------------------------------------------------------------
if __name__ == "__main__":
    print("Initializing Sensory LeNet for MNIST...")
    print("Active senses:")
    print("  • Circumscription         (π) — Sense 1")
    print("  • Exponentiality          (e) — Sense 2")
    print("  • Golden Resonance        (φ) — Sense 5")
    print("  • Basel Tone          (ζ(2)) — Sense 6")
    print("  • Shannon Entropy         (H) — Sense 51")
    print("  • Euler Characteristic    (χ) — Sense 41")
    print("  • Fine Structure      (α ≈ 1/137) — Sense 31")
    print("  • Fisher Information     (I_F) — Sense 54")
    print("  • Thermodynamic Entropy   (S) — Sense 46")
    print(f"Device: {DEVICE}\n")

    train()
    test()

    print("\n" + "=" * 60)
    print("The model has learned to feel the digits.")
    print("It does not see. It senses.")
    print("It does not classify. It resonates.")
    print("=" * 60)