import streamlit as st
import cv2
import numpy as np
import tempfile
from PIL import Image
import torch
from torchvision import transforms, models
import time

# Set page config
st.set_page_config(
    page_title="Dog Language Understanding",
    page_icon="🐕",
    layout="wide"
)

class DogBehaviorAnalyzer:
    def __init__(self):
        # Initialize model (using pretrained ResNet for this example)
        self.model = models.resnet50(pretrained=True)
        self.model.eval()
        
        # Define image transformations
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                              std=[0.229, 0.224, 0.225])
        ])
        
        # Define behavior mappings
        self.behaviors = {
            'tail_wagging': 'Your dog is happy and excited!',
            'barking': 'Your dog is trying to communicate or alert you.',
            'ears_perked': 'Your dog is alert and interested.',
            'lying_down': 'Your dog is relaxed and comfortable.',
            'jumping': 'Your dog is energetic and playful!'
        }

    def analyze_frame(self, frame):
        # Convert frame to PIL Image
        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        
        # Transform image
        input_tensor = self.transform(image)
        input_batch = input_tensor.unsqueeze(0)
        
        # Simulate behavior detection
        # In a real implementation, you'd use a properly trained model
        behaviors = []
        confidence_scores = np.random.random(len(self.behaviors))
        
        for behavior, score in zip(self.behaviors.keys(), confidence_scores):
            if score > 0.7:  # Threshold for detection
                behaviors.append(behavior)
        
        return behaviors

def main():
    st.title("🐕 Dog Language Understanding")
    st.write("Upload a video of your dog to analyze their behavior!")

    # Initialize analyzer
    analyzer = DogBehaviorAnalyzer()

    # File uploader
    video_file = st.file_uploader("Upload Video", type=['mp4', 'avi', 'mov'])

    if video_file is not None:
        # Save uploaded file temporarily
        tfile = tempfile.NamedTemporaryFile(delete=False)
        tfile.write(video_file.read())

        # Video analysis
        cap = cv2.VideoCapture(tfile.name)
        
        # Create columns for layout
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Video Preview")
            video_placeholder = st.empty()
            
        with col2:
            st.subheader("Real-time Analysis")
            analysis_placeholder = st.empty()
            
            # Progress bar
            progress_bar = st.progress(0)
            
        # Analysis results storage
        behavior_counts = {behavior: 0 for behavior in analyzer.behaviors.keys()}
        
        frame_count = 0
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
                
            frame_count += 1
            progress = frame_count / total_frames
            progress_bar.progress(progress)
            
            # Update video preview
            video_placeholder.image(
                cv2.cvtColor(frame, cv2.COLOR_BGR2RGB),
                channels="RGB",
                use_column_width=True
            )
            
            # Analyze frame
            detected_behaviors = analyzer.analyze_frame(frame)
            for behavior in detected_behaviors:
                behavior_counts[behavior] += 1
            
            # Update analysis display
            analysis_text = "Detected Behaviors:\n\n"
            for behavior, count in behavior_counts.items():
                if count > 0:
                    analysis_text += f"• {behavior.replace('_', ' ').title()}: {count} times\n"
                    analysis_text += f"  {analyzer.behaviors[behavior]}\n\n"
            
            analysis_placeholder.text_area(
                "Analysis Results",
                analysis_text,
                height=300
            )
            
            time.sleep(0.1)  # Add small delay for visualization
            
        cap.release()
        
        # Final summary
        st.subheader("Analysis Summary")
        st.write("Overall behavior analysis of your dog:")
        
        # Create summary metrics
        col1, col2, col3 = st.columns(3)
        
        with col1:
            most_common = max(behavior_counts.items(), key=lambda x: x[1])[0]
            st.metric("Most Common Behavior", most_common.replace('_', ' ').title())
            
        with col2:
            total_behaviors = sum(behavior_counts.values())
            st.metric("Total Behaviors Detected", total_behaviors)
            
        with col3:
            behavior_variety = len([b for b in behavior_counts.values() if b > 0])
            st.metric("Behavior Variety", f"{behavior_variety} types")
        
        # Recommendations
        st.subheader("Recommendations")
        if total_behaviors > 0:
            st.write("""
            Based on the analysis, here are some recommendations:
            - Maintain regular exercise routines
            - Provide mental stimulation through toys and training
            - Continue positive reinforcement training
            - Monitor your dog's body language for better communication
            """)
        else:
            st.write("No behaviors detected. Try uploading a different video with clearer dog movements.")

if __name__ == "__main__":
    main()