import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as T
import os
import pathlib
import warnings
import random
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore")

# ==========================================
# 1. Dataset Preparation
# ==========================================

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

base_dir = pathlib.Path('./data/weather_photos/')
class_labels = [p.name for p in base_dir.iterdir() if p.is_dir()]

img_transforms = T.Compose([
    T.Resize((224, 224)),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

dataset = torchvision.datasets.ImageFolder(base_dir, transform=img_transforms)

val_size = int(0.2 * len(dataset))
train_size = len(dataset) - val_size
train_set, val_set = torch.utils.data.random_split(dataset, [train_size, val_size])

loader_params = dict(batch_size=4, shuffle=True, num_workers=1)
train_loader = torch.utils.data.DataLoader(train_set, **loader_params)
val_loader = torch.utils.data.DataLoader(val_set, **loader_params)

# ==========================================
# 2. Core Building Blocks (YOLOv5 Style)
# ==========================================

def smart_pad(kernel, padding=None):
    """Auto padding to maintain spatial dimensions."""
    if padding is None:
        padding = kernel // 2 if isinstance(kernel, int) else [k // 2 for k in kernel]
    return padding

class BaseConv(nn.Module):
    """Standard Convolution -> BatchNorm -> Activation."""
    def __init__(self, in_ch, out_ch, k=1, s=1, p=None, g=1, act=True):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, k, s, smart_pad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(out_ch)
        self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

class ResidualBlock(nn.Module):
    """Standard Bottleneck block."""
    def __init__(self, in_ch, out_ch, use_shortcut=True, g=1, expansion=0.5):
        super().__init__()
        hidden_ch = int(out_ch * expansion)
        self.layer1 = BaseConv(in_ch, hidden_ch, k=1)
        self.layer2 = BaseConv(hidden_ch, out_ch, k=3, g=g)
        self.use_shortcut = use_shortcut and (in_ch == out_ch)

    def forward(self, x):
        residual = x
        x = self.layer2(self.layer1(x))
        if self.use_shortcut:
            x += residual
        return x

class CSPBlock(nn.Module):
    """YOLOv5 C3 Module (Cross Stage Partial Network)."""
    def __init__(self, in_ch, out_ch, n_blocks=1, shortcut=True, g=1, e=0.5):
        super().__init__()
        hidden_ch = int(out_ch * e)
        self.part1 = BaseConv(in_ch, hidden_ch, 1, 1)
        self.part2 = BaseConv(in_ch, hidden_ch, 1, 1)
        self.part3 = BaseConv(2 * hidden_ch, out_ch, 1)
        layers = []
        for _ in range(n_blocks):
            layers.append(ResidualBlock(hidden_ch, hidden_ch, shortcut, g, expansion=1.0))
        self.sub_layers = nn.Sequential(*layers)

    def forward(self, x):
        y1 = self.sub_layers(self.part1(x))
        y2 = self.part2(x)
        return self.part3(torch.cat((y1, y2), dim=1))

# ==========================================
# 3. Classification Model Definition
# ==========================================

class WeatherNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.entry_conv = BaseConv(3, 32, k=3, s=2)
        self.stage1 = CSPBlock(32, 64, n_blocks=3, shortcut=True)
        self.fc_head = nn.Sequential(
            nn.Linear(802816, 100),
            nn.ReLU(),
            nn.Linear(100, 4)
        )

    def forward(self, x):
        x = self.entry_conv(x)
        x = self.stage1(x)
        x = torch.flatten(x, 1)
        return self.fc_head(x)

model = WeatherNet().to(device)
print(f"Model loaded on: {device}")

# ==========================================
# 4. Model Summary
# ==========================================
try:
    import torchsummary
    torchsummary.summary(model, (3, 224, 224))
except ImportError:
    print("Install torchsummary for detailed stats.")

# ==========================================
# 5. Training Utilities
# ==========================================

def run_epoch(dataloader, model, loss_fn, optimizer=None):
    is_training = optimizer is not None
    if is_training:
        model.train()
    else:
        model.eval()

    total_loss, correct = 0, 0
    
    with torch.set_grad_enabled(is_training):
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = loss_fn(outputs, labels)

            if is_training:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            total_loss += loss.item()
            correct += (outputs.argmax(1) == labels).sum().item()

    return correct / len(dataloader.dataset), total_loss / len(dataloader)

# ==========================================
# 6. Training Loop
# ==========================================

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()

history = {'train_acc': [], 'train_loss': [], 'val_acc': [], 'val_loss': []}
best_accuracy = 0.0

for epoch in range(20):
    acc_tr, loss_tr = run_epoch(train_loader, model, criterion, optimizer)
    acc_v, loss_v = run_epoch(val_loader, model, criterion)

    history['train_acc'].append(acc_tr)
    history['train_loss'].append(loss_tr)
    history['val_acc'].append(acc_v)
    history['val_loss'].append(loss_v)

    if acc_v > best_accuracy:
        best_accuracy = acc_v
        torch.save(model.state_dict(), 'best_weather_model.pth')

    print(f"Epoch {epoch+1:02d} | "
          f"Train Acc: {acc_tr*100:.1f}% | Loss: {loss_tr:.3f} || "
          f"Val Acc: {acc_v*100:.1f}% | Loss: {loss_v:.3f}")

# ==========================================
# 7. Visualization
# ==========================================

plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(history['train_acc'], label='Training')
plt.plot(history['val_acc'], label='Validation')
plt.title('Accuracy over Epochs')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(history['train_loss'], label='Training')
plt.plot(history['val_loss'], label='Validation')
plt.title('Loss over Epochs')
plt.legend()

plt.tight_layout()
plt.show()

Implementation Details

Parameter Calculation: The number of learnable parameters in a convolutional layer is calculated as: Kernel_H × Kernel_W × Input_Channels × Output_Channels. For Fully Connected layers, it is Input_Features × Output_Features.

1x1 Convolution Utility:

• Dimensionality Adjustment: Used to increase or decrease the channel depth efficiently.
• Non-linearity Injection: Adds complexity without significant computational overhead.
• Channel Combination: Facilitates information exchange between different channels (feature maps).

Architecture Components:

• BaseConv: A combination of Convolution, Batch Normalization, and SiLU activation. This stacking is standard in modern object detection backbones for stability and speed.
• CSPBlock (C3): Implements the Cross Stage Partial strategy. It splits the feature map, proceses one part through a series of bottlenecks (ResidualBlocks), and concatenates it with the other part before a final integration convolution. This reduces computation while maintaining gradient flow.