import streamlit as st
import torch
import torch.nn as nn
import numpy as np
from huggingface_hub import hf_hub_download

class IcebergClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(2, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2)
        )
        self.fc = nn.Sequential(
            nn.Linear(64 * 9 * 9, 64), nn.ReLU(), 
            nn.Dropout(0.5),
            nn.Linear(64, 1), 
            nn.Sigmoid()
        )

    def forward(self, x): 
        return self.fc(self.conv(x).view(x.size(0), -1))

@st.cache_resource
def load_model():
    model = IcebergClassifier().eval()
    model.load_state_dict(torch.load(hf_hub_download("alperugurcan/iceberg","best_iceberg_model.pth"), map_location='cpu'))
    return model

st.title('🧊 Simple Ship vs Iceberg Detector')

# Simple numeric inputs
band1 = st.number_input('Enter Band 1 value (-40 to -20)', -40.0, -20.0, -30.0)
band2 = st.number_input('Enter Band 2 value (-35 to -15)', -35.0, -15.0, -25.0)

if st.button('Detect'):
    try:
        # Create simple 75x75 arrays with the input values
        b1 = np.full((75,75), band1)
        b2 = np.full((75,75), band2)
        
        # Prepare input tensor
        x = torch.FloatTensor(np.stack([b1,b2])).unsqueeze(0)
        
        # Get prediction
        model = load_model()
        with torch.no_grad():
            pred = model(x).item()
        
        # Show result
        result = "🧊 ICEBERG" if pred > 0.5 else "🚢 SHIP"
        st.success(f"{result} ({pred:.1%})")
        
    except Exception as e:
        st.error(f'Error: {str(e)}')