import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import streamlit as st
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

# Define the Inception Module
class InceptionModule(nn.Module):
    def __init__(self, in_channels):
        super(InceptionModule, self).__init__()
        
        self.branch1x1 = nn.Conv2d(in_channels, 64, kernel_size=1)
        
        self.branch3x3 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
        
        self.branch5x5 = nn.Conv2d(in_channels, 64, kernel_size=5, padding=2)
        
        self.branch_pool = nn.Conv2d(in_channels, 64, kernel_size=1)
        
    def forward(self, x):
        branch1x1 = self.branch1x1(x)
        
        branch3x3 = self.branch3x3(x)
        
        branch5x5 = self.branch5x5(x)
        
        branch_pool = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)(x)
        branch_pool = self.branch_pool(branch_pool)
        
        outputs = [branch1x1, branch3x3, branch5x5, branch_pool]
        return torch.cat(outputs, 1)

# Define the Inception Network
class InceptionNet(nn.Module):
    def __init__(self, num_classes=10):
        super(InceptionNet, self).__init__()
        
        self.inception1 = InceptionModule(in_channels=3)
        self.fc = nn.Linear(64 * 4 * 224 * 224, num_classes)
        
    def forward(self, x):
        x = self.inception1(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Training function
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

# Define a simple transformation
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])

# Initialize model and optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = InceptionNet().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Streamlit app
st.title("InceptionNet Image Classifier")

uploaded_file = st.file_uploader("Choose an image...", type="jpg")
if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image', use_column_width=True)
    
    image = transform(image).unsqueeze(0).to(device)
    
    st.write("Classifying...")
    model.eval()
    with torch.no_grad():
        output = model(image)
    st.write("Output:", output.cpu().detach().numpy())

# Adjust Hyperparameters
st.sidebar.title("Adjust Hyperparameters")
learning_rate = st.sidebar.slider("Learning Rate", 0.0001, 0.01, 0.001)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Visualize Model's Predictions
if st.sidebar.button("Train Model"):
    # Load dummy data for demonstration purposes
    train_loader = DataLoader(
        datasets.FakeData(transform=transform),
        batch_size=64, shuffle=True
    )
    train(model, device, train_loader, optimizer, epoch=1)
    st.sidebar.write("Model Trained!")