"""
Mock model for UI development
This allows parallel development while the actual model is being trained
"""

import torch
import torch.nn as nn
import numpy as np


class MockPlantDiseaseModel(nn.Module):
    """
    Mock CNN model that mimics the structure of a real plant disease classifier
    Returns realistic-looking predictions for UI testing
    """

    def __init__(self, num_classes=39):
        super(MockPlantDiseaseModel, self).__init__()

        # Simple architecture that matches expected input/output
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.AdaptiveAvgPool2d((1, 1))
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64, num_classes)
        )

        self.num_classes = num_classes

    def forward(self, x):
        """
        Forward pass that returns realistic probabilities
        """
        x = self.features(x)
        x = self.classifier(x)

        # Add some controlled randomness to make predictions look realistic
        # In a real model, this would be learned weights
        return x


def create_mock_predictions(class_names):
    """
    Create realistic-looking mock predictions
    Returns a dict with class names and probabilities
    """
    num_classes = len(class_names)

    # Create random probabilities that sum to 1
    # Give higher weight to a few "predicted" classes
    logits = np.random.randn(num_classes)
    logits[np.random.randint(0, num_classes)] += 3  # Make one class likely
    logits[np.random.randint(0, num_classes)] += 1.5  # Make another somewhat likely

    # Convert to probabilities using softmax
    probs = np.exp(logits) / np.sum(np.exp(logits))

    # Create prediction dict
    predictions = {name: float(prob) for name, prob in zip(class_names, probs)}

    return predictions


def get_mock_model():
    """
    Returns a mock model instance
    """
    model = MockPlantDiseaseModel(num_classes=39)
    model.eval()  # Set to evaluation mode
    return model


if __name__ == "__main__":
    # Test the mock model
    print("Testing mock model...")
    model = get_mock_model()

    # Test with random input
    test_input = torch.randn(1, 3, 256, 256)
    with torch.no_grad():
        output = model(test_input)

    print(f"Output shape: {output.shape}")
    print(f"Sample logits: {output[0][:5]}")

    # Test mock predictions
    from config import CLASS_NAMES
    predictions = create_mock_predictions(CLASS_NAMES)
    top_5 = sorted(predictions.items(), key=lambda x: x[1], reverse=True)[:5]

    print("\nTop 5 predictions:")
    for name, prob in top_5:
        print(f"  {name}: {prob:.4f}")