# Complex demo: Load California Housing, preprocess, build/train PyTorch model, evaluate, visualize predictions
import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
import numpy as np
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt

class HousingNet(nn.Module):
    def __init__(self, input_size):
        super(HousingNet, self).__init__()
        self.fc1 = nn.Linear(input_size, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 32)
        self.fc4 = nn.Linear(32, 16)
        self.fc5 = nn.Linear(16, 1)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.relu(self.fc3(x))
        x = self.relu(self.fc4(x))
        x = self.fc5(x)
        return x

def advanced_housing_pytorch():
    # Load and prep
    housing = fetch_california_housing()
    df = pd.DataFrame(housing.data, columns=housing.feature_names)
    df['Target'] = housing.target
    df['RoomsPerHouse'] = df['AveRooms'] / (df['HouseAge'] + 1e-5)  # Feature engineering
    X = df.drop('Target', axis=1).values
    y = df['Target'].values
    # Preprocess
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    # Split
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.25, random_state=42)
    X_train = torch.tensor(X_train, dtype=torch.float32)
    X_test = torch.tensor(X_test, dtype=torch.float32)
    y_train = torch.tensor(y_train, dtype=torch.float32).unsqueeze(1)
    y_test = torch.tensor(y_test, dtype=torch.float32).unsqueeze(1)
    # Build model
    input_size = X_train.shape[1]
    model = HousingNet(input_size)
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    # Train
    losses = []
    for epoch in range(150):
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    # Evaluate
    with torch.no_grad():
        y_pred = model(X_test).numpy().flatten()
        mse = mean_squared_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)
        print(f"Mean Squared Error: {mse:.2f}")
        print(f"R² Score: {r2:.2f}")
    # Visualize predictions
    plt.scatter(y_test, y_pred, c='blue', alpha=0.5)
    plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=2)
    plt.xlabel('True Values')
    plt.ylabel('Predictions')
    plt.title('California Housing PyTorch Predictions')
    plt.annotate('Best Fit', xy=(y_test.mean(), y_pred.mean()), xytext=(y_test.mean() + 0.5, y_pred.mean() + 0.5), arrowprops=dict(facecolor='black', shrink=0.05))
    plt.show()
    # Visualize loss
    plt.plot(losses, label='Training Loss')
    plt.title('Housing PyTorch Model Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.annotate('Convergence', xy=(120, losses[-1]), xytext=(100, losses[-1] + 0.1), arrowprops=dict(facecolor='black', shrink=0.05))
    plt.show()

advanced_housing_pytorch()