# Complex demo: Load Iris, preprocess, build/train PyTorch model, evaluate, visualize loss and confusion matrix
import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns

class IrisNet(nn.Module):
    def __init__(self):
        super(IrisNet, self).__init__()
        self.fc1 = nn.Linear(4, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 16)
        self.fc4 = nn.Linear(16, 3)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)

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

def advanced_iris_pytorch():
    # Load and prep
    iris = load_iris()
    df = pd.DataFrame(iris.data, columns=iris.feature_names)
    df['Target'] = iris.target
    df['Petal_Ratio'] = df['petal length (cm)'] / (df['petal width (cm)'] + 1e-5)  # Feature engineering
    X = df.drop('Target', axis=1).values
    y = df['Target'].values
    # Preprocess
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.3, random_state=42, stratify=y)
    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.long)
    y_test = torch.tensor(y_test, dtype=torch.long)
    # Build model
    model = IrisNet()
    criterion = nn.CrossEntropyLoss()
    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)
        y_pred_classes = torch.argmax(y_pred, dim=1).numpy()
        acc = accuracy_score(y_test, y_pred_classes)
        cm = confusion_matrix(y_test, y_pred_classes)
        print(f"Accuracy: {acc:.2f}")
        print("Confusion Matrix:\n", cm)
    # Visualize loss
    plt.plot(losses, label='Training Loss')
    plt.title('Iris 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()
    # Visualize confusion matrix
    sns.heatmap(cm, annot=True, cmap='Blues', fmt='d')
    plt.title('Confusion Matrix for Iris PyTorch')
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.annotate('High Accuracy', xy=(1.5, 0.5), xytext=(2.5, 1.5), arrowprops=dict(facecolor='black', shrink=0.05))
    plt.show()

advanced_iris_pytorch()