import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
import pandas as pd
from statsmodels.tsa.stattools import grangercausalitytests

class GravityNet(nn.Module):
    def __init__(self):
        super(GravityNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(16 * 14 * 14, 10)
        )
    def forward(self, x):
        return self.features(x)

def apply_missing_info_gravity(model, base_strength, test_accuracy):
    """
    Gravity is proportional to the missing information (1 - Test Accuracy).
    If accuracy is low, gravity pulls heavily. If accuracy is 100%, gravity becomes 0.
    """
    # The 'missing information' coefficient
    missing_info_factor = 1.0 - test_accuracy  
    dynamic_strength = base_strength * missing_info_factor
    
    with torch.no_grad():
        for param in model.parameters():
            if param.grad is not None:
                param.grad.add_(dynamic_strength * torch.sign(param.data))
    return dynamic_strength

def run_experiment(use_gravity, base_gravity_strength=0.1):
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_loader = DataLoader(datasets.MNIST('../data', train=True, download=True, transform=transform), batch_size=128, shuffle=True)
    test_loader = DataLoader(datasets.MNIST('../data', train=False, transform=transform), batch_size=256, shuffle=True)
    
    # Create an iterator for the test set to sample "missing information" on the fly
    test_iter = iter(test_loader)
    
    model = GravityNet()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    criterion = nn.CrossEntropyLoss()
    
    accuracy_history = []
    
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        
        # Initialize a running tracking metric for missing information
        ema_test_acc = 0.10  # Start with an assumption of random guess (10% for MNIST)
        alpha = 0.2          # Smoothing factor (how much weight to give the new batch)
        
        # Inside your training loop:
        if use_gravity:
            try:
                test_data, test_target = next(test_iter)
            except StopIteration:
                test_iter = iter(test_loader)
                test_data, test_target = next(test_iter)
            
            with torch.no_grad():
                test_output = model(test_data)
                pred = test_output.argmax(dim=1, keepdim=True)
                batch_acc = pred.eq(test_target.view_as(pred)).sum().item() / len(test_data)
            
            # Smooth out the vector field's fluctuations
            ema_test_acc = (alpha * batch_acc) + ((1 - alpha) * ema_test_acc)
            
            # Apply the stable gravitational pull
            apply_missing_info_gravity(model, base_gravity_strength, ema_test_acc)
                    
        optimizer.step()
        
        if batch_idx % 5 == 0:
            pred = output.argmax(dim=1, keepdim=True)
            acc = pred.eq(target.view_as(pred)).sum().item() / len(data)
            accuracy_history.append(acc)
            
        if batch_idx >= 150: 
            break
            
    return accuracy_history

# --- Causality Testing Matrix ---
print("Running Baseline Control...")
control_history = run_experiment(use_gravity=False)

print("Running Missing Information Gravity Experiment...")
gravity_history = run_experiment(use_gravity=True, base_gravity_strength=0.15)

# Build a joint dataframe to see if one history forces the other
df = pd.DataFrame({
    'Control': control_history,
    'Gravity': gravity_history
}).diff().dropna()  # Stationarize trends

print("\n--- Testing Environmental Causality ---")
try:
    # We test if the Gravity trajectory Granger-causes the Control trajectory's variance
    gc_res = grangercausalitytests(df[['Control', 'Gravity']], maxlag=[2], verbose=False)
    p_value = gc_res[2][0]['ssr_ftest'][1]
    print(f"Granger Causality p-value: {p_value:.5f}")
    if p_value < 0.05:
        print("Verdict: SUCCESS. Tying gravity to missing information creates a detectable causal imprint on the vector field.")
    else:
        print("Verdict: p-value still high. The model might need more training steps/epochs to establish a steady directional flow.")
except Exception as e:
    print(f"Causality calculation skipped: {e}")