#!/usr/bin/env python3
"""
Generate boilerplate code for common GNN architectures in PyTorch Geometric.

This script creates ready-to-use GNN model templates with training loops,
evaluation metrics, and proper data handling.

Usage:
    python create_gnn_template.py --model gcn --task node_classification --output my_model.py
    python create_gnn_template.py --model gat --task graph_classification --output graph_classifier.py
"""

import argparse
from pathlib import Path


TEMPLATES = {
    'node_classification': {
        'gcn': '''import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.datasets import Planetoid


class GCN(torch.nn.Module):
    """Graph Convolutional Network for node classification."""

    def __init__(self, num_features, hidden_channels, num_classes, num_layers=2, dropout=0.5):
        super().__init__()
        self.convs = torch.nn.ModuleList()

        # First layer
        self.convs.append(GCNConv(num_features, hidden_channels))

        # Hidden layers
        for _ in range(num_layers - 2):
            self.convs.append(GCNConv(hidden_channels, hidden_channels))

        # Output layer
        self.convs.append(GCNConv(hidden_channels, num_classes))

        self.dropout = dropout

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        # Apply conv layers with ReLU and dropout
        for conv in self.convs[:-1]:
            x = conv(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)

        # Final layer without activation
        x = self.convs[-1](x, edge_index)
        return F.log_softmax(x, dim=1)


def train(model, data, optimizer):
    """Train the model for one epoch."""
    model.train()
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()


@torch.no_grad()
def test(model, data):
    """Evaluate the model."""
    model.eval()
    out = model(data)
    pred = out.argmax(dim=1)

    accs = []
    for mask in [data.train_mask, data.val_mask, data.test_mask]:
        correct = (pred[mask] == data.y[mask]).sum()
        accs.append(int(correct) / int(mask.sum()))

    return accs


def main():
    # Load dataset
    dataset = Planetoid(root='/tmp/Cora', name='Cora')
    data = dataset[0]

    # Create model
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = GCN(
        num_features=dataset.num_features,
        hidden_channels=64,
        num_classes=dataset.num_classes,
        num_layers=3,
        dropout=0.5
    ).to(device)
    data = data.to(device)

    # Setup optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

    # Training loop
    print("Training GCN model...")
    best_val_acc = 0
    for epoch in range(1, 201):
        loss = train(model, data, optimizer)
        train_acc, val_acc, test_acc = test(model, data)

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_test_acc = test_acc

        if epoch % 10 == 0:
            print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
                  f'Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}')

    print(f'\\nBest Test Accuracy: {best_test_acc:.4f}')


if __name__ == '__main__':
    main()
''',

        'gat': '''import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv
from torch_geometric.datasets import Planetoid


class GAT(torch.nn.Module):
    """Graph Attention Network for node classification."""

    def __init__(self, num_features, hidden_channels, num_classes, heads=8, dropout=0.6):
        super().__init__()

        self.conv1 = GATConv(num_features, hidden_channels, heads=heads, dropout=dropout)
        self.conv2 = GATConv(hidden_channels * heads, num_classes, heads=1,
                             concat=False, dropout=dropout)

        self.dropout = dropout

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = F.dropout(x, p=self.dropout, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)


def train(model, data, optimizer):
    """Train the model for one epoch."""
    model.train()
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()


@torch.no_grad()
def test(model, data):
    """Evaluate the model."""
    model.eval()
    out = model(data)
    pred = out.argmax(dim=1)

    accs = []
    for mask in [data.train_mask, data.val_mask, data.test_mask]:
        correct = (pred[mask] == data.y[mask]).sum()
        accs.append(int(correct) / int(mask.sum()))

    return accs


def main():
    # Load dataset
    dataset = Planetoid(root='/tmp/Cora', name='Cora')
    data = dataset[0]

    # Create model
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = GAT(
        num_features=dataset.num_features,
        hidden_channels=8,
        num_classes=dataset.num_classes,
        heads=8,
        dropout=0.6
    ).to(device)
    data = data.to(device)

    # Setup optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

    # Training loop
    print("Training GAT model...")
    best_val_acc = 0
    for epoch in range(1, 201):
        loss = train(model, data, optimizer)
        train_acc, val_acc, test_acc = test(model, data)

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_test_acc = test_acc

        if epoch % 10 == 0:
            print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
                  f'Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}')

    print(f'\\nBest Test Accuracy: {best_test_acc:.4f}')


if __name__ == '__main__':
    main()
''',

        'graphsage': '''import torch
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from torch_geometric.datasets import Planetoid


class GraphSAGE(torch.nn.Module):
    """GraphSAGE for node classification."""

    def __init__(self, num_features, hidden_channels, num_classes, num_layers=2, dropout=0.5):
        super().__init__()
        self.convs = torch.nn.ModuleList()

        self.convs.append(SAGEConv(num_features, hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(SAGEConv(hidden_channels, hidden_channels))
        self.convs.append(SAGEConv(hidden_channels, num_classes))

        self.dropout = dropout

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        for conv in self.convs[:-1]:
            x = conv(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)

        x = self.convs[-1](x, edge_index)
        return F.log_softmax(x, dim=1)


def train(model, data, optimizer):
    model.train()
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()


@torch.no_grad()
def test(model, data):
    model.eval()
    out = model(data)
    pred = out.argmax(dim=1)

    accs = []
    for mask in [data.train_mask, data.val_mask, data.test_mask]:
        correct = (pred[mask] == data.y[mask]).sum()
        accs.append(int(correct) / int(mask.sum()))

    return accs


def main():
    dataset = Planetoid(root='/tmp/Cora', name='Cora')
    data = dataset[0]

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = GraphSAGE(
        num_features=dataset.num_features,
        hidden_channels=64,
        num_classes=dataset.num_classes,
        num_layers=2,
        dropout=0.5
    ).to(device)
    data = data.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

    print("Training GraphSAGE model...")
    best_val_acc = 0
    for epoch in range(1, 201):
        loss = train(model, data, optimizer)
        train_acc, val_acc, test_acc = test(model, data)

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_test_acc = test_acc

        if epoch % 10 == 0:
            print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
                  f'Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}')

    print(f'\\nBest Test Accuracy: {best_test_acc:.4f}')


if __name__ == '__main__':
    main()
''',
    },

    'graph_classification': {
        'gin': '''import torch
import torch.nn.functional as F
from torch_geometric.nn import GINConv, global_add_pool
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader


class GIN(torch.nn.Module):
    """Graph Isomorphism Network for graph classification."""

    def __init__(self, num_features, hidden_channels, num_classes, num_layers=3, dropout=0.5):
        super().__init__()

        self.convs = torch.nn.ModuleList()
        self.batch_norms = torch.nn.ModuleList()

        # Create MLP for first layer
        nn = torch.nn.Sequential(
            torch.nn.Linear(num_features, hidden_channels),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_channels, hidden_channels)
        )
        self.convs.append(GINConv(nn))
        self.batch_norms.append(torch.nn.BatchNorm1d(hidden_channels))

        # Hidden layers
        for _ in range(num_layers - 2):
            nn = torch.nn.Sequential(
                torch.nn.Linear(hidden_channels, hidden_channels),
                torch.nn.ReLU(),
                torch.nn.Linear(hidden_channels, hidden_channels)
            )
            self.convs.append(GINConv(nn))
            self.batch_norms.append(torch.nn.BatchNorm1d(hidden_channels))

        # Output MLP
        self.lin = torch.nn.Linear(hidden_channels, num_classes)
        self.dropout = dropout

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch

        for conv, batch_norm in zip(self.convs, self.batch_norms):
            x = conv(x, edge_index)
            x = batch_norm(x)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)

        # Global pooling
        x = global_add_pool(x, batch)

        # Output layer
        x = self.lin(x)
        return F.log_softmax(x, dim=1)


def train(model, loader, optimizer, device):
    """Train the model for one epoch."""
    model.train()
    total_loss = 0

    for data in loader:
        data = data.to(device)
        optimizer.zero_grad()
        out = model(data)
        loss = F.nll_loss(out, data.y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs

    return total_loss / len(loader.dataset)


@torch.no_grad()
def test(model, loader, device):
    """Evaluate the model."""
    model.eval()
    correct = 0

    for data in loader:
        data = data.to(device)
        out = model(data)
        pred = out.argmax(dim=1)
        correct += (pred == data.y).sum().item()

    return correct / len(loader.dataset)


def main():
    # Load dataset
    dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
    print(f"Dataset: {dataset}")
    print(f"Number of graphs: {len(dataset)}")
    print(f"Number of features: {dataset.num_features}")
    print(f"Number of classes: {dataset.num_classes}")

    # Shuffle and split
    dataset = dataset.shuffle()
    train_dataset = dataset[:int(len(dataset) * 0.8)]
    test_dataset = dataset[int(len(dataset) * 0.8):]

    # Create data loaders
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=32)

    # Create model
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = GIN(
        num_features=dataset.num_features,
        hidden_channels=64,
        num_classes=dataset.num_classes,
        num_layers=3,
        dropout=0.5
    ).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    # Training loop
    print("\\nTraining GIN model...")
    for epoch in range(1, 101):
        loss = train(model, train_loader, optimizer, device)
        train_acc = test(model, train_loader, device)
        test_acc = test(model, test_loader, device)

        if epoch % 10 == 0:
            print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
                  f'Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}')


if __name__ == '__main__':
    main()
''',
    },
}


def generate_template(model_type: str, task: str, output_path: str):
    """Generate a GNN template file."""
    if task not in TEMPLATES:
        raise ValueError(f"Unknown task: {task}. Available: {list(TEMPLATES.keys())}")

    if model_type not in TEMPLATES[task]:
        raise ValueError(f"Model {model_type} not available for task {task}. "
                         f"Available: {list(TEMPLATES[task].keys())}")

    template = TEMPLATES[task][model_type]

    # Write to file
    output_file = Path(output_path)
    output_file.parent.mkdir(parents=True, exist_ok=True)

    with open(output_file, 'w') as f:
        f.write(template)

    print(f"✓ Generated {model_type.upper()} template for {task}")
    print(f"  Saved to: {output_path}")
    print(f"\\nTo run the template:")
    print(f"  python {output_path}")


def list_templates():
    """List all available templates."""
    print("Available GNN Templates")
    print("=" * 50)
    for task, models in TEMPLATES.items():
        print(f"\\n{task.upper()}")
        print("-" * 50)
        for model in models.keys():
            print(f"  - {model}")
    print()


def main():
    parser = argparse.ArgumentParser(
        description="Generate GNN model templates",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python create_gnn_template.py --model gcn --task node_classification --output gcn_model.py
  python create_gnn_template.py --model gin --task graph_classification --output gin_model.py
  python create_gnn_template.py --list
        """
    )

    parser.add_argument('--model', type=str,
                        help='Model type (gcn, gat, graphsage, gin)')
    parser.add_argument('--task', type=str,
                        help='Task type (node_classification, graph_classification)')
    parser.add_argument('--output', type=str, default='gnn_model.py',
                        help='Output file path (default: gnn_model.py)')
    parser.add_argument('--list', action='store_true',
                        help='List all available templates')

    args = parser.parse_args()

    if args.list:
        list_templates()
        return

    if not args.model or not args.task:
        parser.print_help()
        print("\\n" + "=" * 50)
        list_templates()
        return

    try:
        generate_template(args.model, args.task, args.output)
    except ValueError as e:
        print(f"Error: {e}")
        print("\\nUse --list to see available templates")


if __name__ == '__main__':
    main()