"""
Travel Planning System with Multi-Agent Architecture - Gradio Interface

This Gradio app provides an interactive interface for the travel planning system
with multiple specialized agents for weather analysis, hotel recommendations,
itinerary planning, and trip summarization.
"""

import gradio as gr
import numpy as np
from typing import List, Dict
from sklearn.ensemble import RandomForestRegressor
from sentence_transformers import SentenceTransformer
from transformers import pipeline
import warnings
warnings.filterwarnings("ignore")


class WeatherAnalysisAgent:
    """An agent that analyzes historical weather data to predict optimal travel periods."""
    def __init__(self):
        self.model = RandomForestRegressor(n_estimators=100)

    def train(self, historical_data: Dict):
        X = np.array([[d['month'], d['latitude'], d['longitude']] for d in historical_data])
        y = np.array([d['weather_score'] for d in historical_data])
        self.model.fit(X, y)

    def predict_best_time(self, location: Dict) -> Dict:
        predictions = []
        for month in range(1, 13):
            prediction = self.model.predict([[
                month,
                location['latitude'],
                location['longitude']
            ]]).item()
            predictions.append({'month': month, 'score': float(prediction)})

        return {
            'best_months': sorted(predictions, key=lambda x: x['score'], reverse=True)[:3],
            'location': location
        }


class HotelRecommenderAgent:
    """An agent for recommending hotels using semantic embeddings."""
    def __init__(self):
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
        self.hotels_db = []
        self.hotels_embeddings = None

    def add_hotels(self, hotels: List[Dict]):
        self.hotels_db = hotels
        descriptions = [h['description'] for h in hotels]
        self.hotels_embeddings = self.encoder.encode(descriptions)

    def find_hotels(self, preferences: str, top_k: int = 5) -> List[Dict]:
        pref_embedding = self.encoder.encode([preferences])
        similarities = np.dot(self.hotels_embeddings, pref_embedding.T).flatten()

        top_indices = similarities.argsort()[-top_k:][::-1]
        return [
            {**self.hotels_db[i], 'similarity_score': float(similarities[i])}
            for i in top_indices
        ]


class ItineraryPlannerAgent:
    """An agent that creates travel itineraries using text generation."""
    def __init__(self):
        self.planner = pipeline(
            "text-generation",
            model="gpt2",
            max_new_tokens=200,
            truncation=True,
            pad_token_id=50256
        )

    def create_itinerary(self, destination_info: Dict, weather_info: Dict,
                         hotel_info: Dict, duration: int) -> Dict:
        prompt = self._create_prompt(destination_info, weather_info, hotel_info, duration)
        response = self.planner(prompt)[0]['generated_text']

        return {
            'itinerary': response,
            'duration': duration,
            'destination': destination_info['name']
        }

    def _create_prompt(self, destination_info: Dict, weather_info: Dict,
                       hotel_info: Dict, duration: int) -> str:
        return f"Day-by-day itinerary for {destination_info['name']} ({duration} days):"


class SummaryAgent:
    """An agent that generates summarized travel plans."""
    def __init__(self):
        self.llm = pipeline(
            "text-generation",
            model="gpt2",
            max_new_tokens=300,
            truncation=True,
            pad_token_id=50256
        )

    def calculate_total_price(self, hotel_info: Dict, duration: int) -> float:
        hotel_cost = hotel_info[0]['price'] * duration
        daily_expenses = 100
        additional_costs = daily_expenses * duration
        return hotel_cost + additional_costs

    def create_email(self, trip_data: Dict, client_name: str) -> Dict:
        total_price = self.calculate_total_price(
            trip_data['recommended_hotels'],
            trip_data['itinerary']['duration']
        )

        prompt = f"Travel summary for {client_name}: {trip_data['itinerary']['destination']} trip"

        response = self.llm(prompt)[0]['generated_text']

        return {
            'email_content': response,
            'total_price': total_price,
            'summary_data': {
                'destination': trip_data['itinerary']['destination'],
                'duration': trip_data['itinerary']['duration'],
                'hotel': trip_data['recommended_hotels'][0]['name'],
                'best_month': trip_data['weather_analysis']['best_months'][0]['month']
            }
        }


class TravelPlanningSystem:
    def __init__(self):
        self.weather_agent = WeatherAnalysisAgent()
        self.hotel_agent = HotelRecommenderAgent()
        self.itinerary_agent = ItineraryPlannerAgent()
        self.summary_agent = SummaryAgent()
        self.initialized = False

    def setup(self, historical_weather_data: Dict, hotels_database: List[Dict]):
        try:
            self.weather_agent.train(historical_weather_data)
            self.hotel_agent.add_hotels(hotels_database)
            self.initialized = True
        except Exception as e:
            print(f"Setup error: {e}")
            self.initialized = False

    def plan_trip(self, destination: Dict, preferences: str, duration: int, client_name: str) -> Dict:
        if not self.initialized:
            return {"error": "System not properly initialized"}

        try:
            # 1. Weather analysis
            weather_analysis = self.weather_agent.predict_best_time(destination)

            # 2. Hotel search
            recommended_hotels = self.hotel_agent.find_hotels(preferences)

            # 3. Itinerary creation
            itinerary = self.itinerary_agent.create_itinerary(
                destination,
                weather_analysis,
                recommended_hotels,
                duration
            )

            # 4. Create summary
            trip_data = {
                'weather_analysis': weather_analysis,
                'recommended_hotels': recommended_hotels,
                'itinerary': itinerary
            }

            summary = self.summary_agent.create_email(trip_data, client_name)

            return {
                **trip_data,
                'summary': summary
            }
        except Exception as e:
            return {"error": f"Planning error: {str(e)}"}


# Initialize the system
travel_system = TravelPlanningSystem()

# Sample data
historical_weather_data = [
    {'month': 1, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.5},
    {'month': 2, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.6},
    {'month': 3, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.7},
    {'month': 4, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.8},
    {'month': 5, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.85},
    {'month': 6, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.9},
    {'month': 7, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.95},
    {'month': 8, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.9},
    {'month': 9, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.85},
    {'month': 10, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.7},
    {'month': 11, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.6},
    {'month': 12, 'latitude': 41.9028, 'longitude': 12.4964, 'weather_score': 0.5},
]

hotels_database = [
    {
        'name': 'Grand Hotel',
        'description': 'Luxury hotel in city center with spa and restaurant',
        'price': 300
    },
    {
        'name': 'Boutique Resort',
        'description': 'Intimate boutique hotel with personalized service',
        'price': 250
    },
    {
        'name': 'City View Hotel',
        'description': 'Modern hotel with panoramic city views',
        'price': 200
    }
]

# Initialize the system
travel_system.setup(historical_weather_data, hotels_database)


def plan_trip_interface(destination_name, latitude, longitude, attractions, preferences, duration, client_name):
    """Gradio interface function for trip planning."""
    try:
        # Validate inputs
        if not destination_name or not destination_name.strip():
            return "❌ Please enter a destination name."
        
        if not client_name or not client_name.strip():
            return "❌ Please enter your name."
        
        if not preferences or not preferences.strip():
            return "❌ Please enter your hotel preferences."
        
        if not attractions or not attractions.strip():
            return "❌ Please enter at least one attraction."
        
        if duration <= 0:
            return "❌ Trip duration must be at least 1 day."
        
        # Convert string inputs to appropriate types
        lat = float(latitude)
        lon = float(longitude)
        dur = int(duration)
        
        # Validate coordinates
        if not (-90 <= lat <= 90):
            return "❌ Latitude must be between -90 and 90."
        if not (-180 <= lon <= 180):
            return "❌ Longitude must be between -180 and 180."
        
        # Parse attractions (comma-separated string)
        attractions_list = [attr.strip() for attr in attractions.split(',') if attr.strip()]
        
        if not attractions_list:
            return "❌ Please enter at least one valid attraction."
        
        destination = {
            'name': destination_name.strip(),
            'latitude': lat,
            'longitude': lon,
            'attractions': attractions_list
        }
        
        # Plan the trip
        result = travel_system.plan_trip(destination, preferences.strip(), dur, client_name.strip())
        
        if "error" in result:
            return f"❌ **Error:** {result['error']}"
        
        # Format the output
        output = f"""
# 🌍 Travel Plan for {client_name.strip()}

## 📍 Destination: {destination_name.strip()}
**Duration:** {dur} days  
**Best time to visit:** Month {result['weather_analysis']['best_months'][0]['month']} (Score: {result['weather_analysis']['best_months'][0]['score']:.2f})

## 🏨 Recommended Hotels:
"""
        for i, hotel in enumerate(result['recommended_hotels'][:3], 1):
            output += f"{i}. **{hotel['name']}** - ${hotel['price']}/night (Match: {hotel['similarity_score']:.2f})\n"
            output += f"   _{hotel['description']}_\n\n"

        output += f"## 📅 Itinerary:\n{result['itinerary']['itinerary']}\n\n"
        output += f"## 💰 Estimated Total Cost: **${result['summary']['total_price']:.2f}**\n\n"
        output += f"## 📧 Email Summary:\n{result['summary']['email_content']}"
        
        return output
        
    except ValueError as e:
        return f"❌ **Invalid input:** {str(e)}"
    except Exception as e:
        return f"❌ **Unexpected error:** {str(e)}"


# Create Gradio interface
with gr.Blocks(title="Travel Planning System", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌍 AI Travel Planning System")
    gr.Markdown("Plan your perfect trip with our multi-agent AI system that analyzes weather, recommends hotels, and creates detailed itineraries.")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## 📝 Trip Details")
            destination_name = gr.Textbox(
                label="Destination Name", 
                placeholder="e.g., Rome, Paris, Tokyo",
                info="Enter the city or country you want to visit"
            )
            latitude = gr.Number(
                label="Latitude", 
                value=41.9028, 
                precision=4,
                info="GPS latitude coordinate (-90 to 90)"
            )
            longitude = gr.Number(
                label="Longitude", 
                value=12.4964, 
                precision=4,
                info="GPS longitude coordinate (-180 to 180)"
            )
            attractions = gr.Textbox(
                label="Attractions (comma-separated)", 
                placeholder="e.g., Colosseum, Vatican, Trevi Fountain",
                info="List the places you want to visit"
            )
            preferences = gr.Textbox(
                label="Hotel Preferences", 
                placeholder="e.g., Luxury hotel with spa facilities", 
                lines=3,
                info="Describe your ideal accommodation"
            )
            duration = gr.Number(
                label="Trip Duration (days)", 
                value=3, 
                precision=0,
                minimum=1,
                maximum=30,
                info="How many days will you stay?"
            )
            client_name = gr.Textbox(
                label="Your Name", 
                placeholder="e.g., John Smith",
                info="For personalized recommendations"
            )
            
            plan_btn = gr.Button("🚀 Plan My Trip", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            gr.Markdown("## 📋 Your Travel Plan")
            output = gr.Markdown(
                value="👆 Fill in the details on the left and click 'Plan My Trip' to get started!"
            )
    
    # Event handler with loading state
    plan_btn.click(
        fn=plan_trip_interface,
        inputs=[destination_name, latitude, longitude, attractions, preferences, duration, client_name],
        outputs=output,
        show_progress=True
    )
    
    # Example section
    with gr.Accordion("💡 Try these examples", open=False):
        gr.Examples(
            examples=[
                ["Rome", 41.9028, 12.4964, "Colosseum, Vatican, Trevi Fountain", "Luxury hotel with spa facilities", 3, "John Smith"],
                ["Paris", 48.8566, 2.3522, "Eiffel Tower, Louvre, Notre-Dame", "Boutique hotel near attractions", 5, "Jane Doe"],
                ["Tokyo", 35.6762, 139.6503, "Senso-ji Temple, Tokyo Skytree, Meiji Shrine", "Modern hotel with city views", 4, "Mike Johnson"]
            ],
            inputs=[destination_name, latitude, longitude, attractions, preferences, duration, client_name]
        )
    
    # Footer
    gr.Markdown("---")
    gr.Markdown("""
    ### 🤖 About this AI System
    This travel planning system uses multiple AI agents:
    - **Weather Agent**: Analyzes historical weather data to recommend the best travel months
    - **Hotel Agent**: Uses semantic similarity to match your preferences with hotel descriptions  
    - **Itinerary Agent**: Generates detailed day-by-day travel plans
    - **Summary Agent**: Creates personalized trip summaries and cost estimates
    """)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)