import torch
from torchvision import models, transforms
from PIL import Image

class PlantDiseaseClassifier:
    """
    A class to handle plant disease classification using a pre-trained model.
    This is a placeholder for the actual model implementation.
    """
    def __init__(self, model_path=None, class_names=None):
        # Set up the model architecture
        self.model = models.resnet50(pretrained=True)
        num_ftrs = self.model.fc.in_features
        
        # If we have a specific number of classes, replace the final layer
        if class_names:
            self.model.fc = torch.nn.Linear(num_ftrs, len(class_names))
            self.class_names = class_names
        else:
            # Default class names (placeholder)
            self.class_names = [
                "Apple___Apple_scab", 
                "Apple___Black_rot", 
                "Apple___Cedar_apple_rust", 
                "Apple___healthy", 
                "Corn_(maize)___Cercospora_leaf_spot", 
                "Corn_(maize)___Common_rust", 
                "Corn_(maize)___healthy", 
                "Tomato___Early_blight", 
                "Tomato___Late_blight", 
                "Tomato___healthy"
            ]
            self.model.fc = torch.nn.Linear(num_ftrs, len(self.class_names))
        
        # Load pretrained weights if provided
        if model_path:
            self.model.load_state_dict(torch.load(model_path))
        
        # Set up image transformation
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        # Set model to evaluation mode
        self.model.eval()
    
    def predict(self, image):
        """
        Predict the disease class for a given plant image
        
        Args:
            image: PIL Image of the plant
            
        Returns:
            class_name: String representing the predicted disease class
            confidence: Float representing the confidence score
        """
        # Transform the image
        img_tensor = self.transform(image).unsqueeze(0)
        
        # Get prediction
        with torch.no_grad():
            outputs = self.model(img_tensor)
            _, predictions = torch.max(outputs, 1)
            confidence = torch.nn.functional.softmax(outputs, dim=1)[0][predictions.item()].item()
        
        # Get class name
        class_name = self.class_names[predictions.item()]
        
        return class_name, confidence

class QuestionAnswerer:
    """
    A class to handle question answering about plant diseases.
    This is a placeholder for more sophisticated NLP models.
    """
    def __init__(self, treatments_df):
        self.treatments_df = treatments_df
    
    def answer(self, question):
        """
        Answer a question about plant diseases using the treatments dataframe
        
        Args:
            question: String representing the user's question
            
        Returns:
            answer: String representing the answer to the question
        """
        question_lower = question.lower()
        
        # Extract potential crop and disease from question
        crops = list(set(self.treatments_df['Crop'].str.lower()))
        diseases = list(set(self.treatments_df['Disease'].str.lower()))
        
        # Find mentions of crops
        mentioned_crops = [crop for crop in crops if crop in question_lower]
        
        # Find mentions of diseases
        mentioned_diseases = [disease for disease in diseases if disease in question_lower]
        
        # Check question intent
        if any(term in question_lower for term in ["treat", "treatment", "how to", "cure"]):
            # Question about treatment
            if mentioned_crops and mentioned_diseases:
                # Specific crop and disease
                matches = self.treatments_df[
                    (self.treatments_df['Crop'].str.lower().isin(mentioned_crops)) & 
                    (self.treatments_df['Disease'].str.lower().isin(mentioned_diseases))
                ]
                if not matches.empty:
                    match = matches.iloc[0]
                    return f"Treatment for {match['Disease']} in {match['Crop']}: {match['Treatment']}\n\nChemical control: {match['Medicine/Chemical Control']}"
            
            # Just disease mentioned
            elif mentioned_diseases:
                matches = self.treatments_df[self.treatments_df['Disease'].str.lower().isin(mentioned_diseases)]
                if not matches.empty:
                    match = matches.iloc[0]
                    return f"Treatment for {match['Disease']}: {match['Treatment']}\n\nChemical control: {match['Medicine/Chemical Control']}"
        
        elif any(term in question_lower for term in ["symptom", "sign", "identify"]):
            # Question about symptoms
            if mentioned_diseases:
                matches = self.treatments_df[self.treatments_df['Disease'].str.lower().isin(mentioned_diseases)]
                if not matches.empty:
                    match = matches.iloc[0]
                    return f"Symptoms of {match['Disease']}: {match['Symptoms']}"
        
        # Generic response
        return "I couldn't find specific information for your question. Please try uploading an image of the plant or asking about a specific crop disease treatment."