import os
import json
import numpy as np
from datetime import datetime, timedelta
import time
from PIL import Image
import cv2
from geopy.distance import geodesic
import math

# Optional advanced dependencies with fallback handling
try:
    import pandas as pd
    import geopandas as gpd
    PANDAS_AVAILABLE = True
except ImportError:
    PANDAS_AVAILABLE = False
    print("Warning: pandas/geopandas not available. Install with: pip install pandas geopandas")

try:
    import plotly.graph_objects as go
    import plotly.express as px
    from plotly.subplots import make_subplots
    PLOTLY_AVAILABLE = True
except ImportError:
    PLOTLY_AVAILABLE = False
    print("Warning: plotly not available. Install with: pip install plotly")

try:
    import pydeck as pdk
    PYDECK_AVAILABLE = True
except ImportError:
    PYDECK_AVAILABLE = False
    print("Warning: pydeck not available. Install with: pip install pydeck")

try:
    from moviepy import VideoFileClip, ImageSequenceClip
    MOVIEPY_AVAILABLE = True
except ImportError:
    MOVIEPY_AVAILABLE = False
    print("Warning: moviepy not available. Install with: pip install moviepy")

class Advanced3DGenerator:
    """
    Advanced 3D animation generator using Pydeck + Plotly + Blender pipeline
    for high-quality GPS track visualizations
    """
    
    def __init__(self, output_folder):
        self.output_folder = output_folder
        self.frames_folder = os.path.join(output_folder, "frames")
        self.temp_folder = os.path.join(output_folder, "temp")
        
        # Create necessary folders
        os.makedirs(self.frames_folder, exist_ok=True)
        os.makedirs(self.temp_folder, exist_ok=True)
        
        # Animation settings
        self.fps = 30
        self.duration_per_point = 0.5  # seconds per GPS point
        self.camera_height = 1000  # meters
        self.trail_length = 50  # number of previous points to show
        
    def check_dependencies(self):
        """Check if all required dependencies are available"""
        missing = []
        if not PANDAS_AVAILABLE:
            missing.append("pandas, geopandas")
        if not PLOTLY_AVAILABLE:
            missing.append("plotly")
        if not PYDECK_AVAILABLE:
            missing.append("pydeck")
        if not MOVIEPY_AVAILABLE:
            missing.append("moviepy")
            
        if missing:
            raise ImportError(f"Missing required dependencies: {', '.join(missing)}. Please install them with: pip install {' '.join(missing)}")
        
        return True
        
    def load_gps_data(self, positions_file):
        """Load and preprocess GPS data"""
        with open(positions_file, 'r') as f:
            positions = json.load(f)
        
        # Convert to DataFrame
        df = pd.DataFrame(positions)
        
        # Parse timestamps
        df['timestamp'] = pd.to_datetime(df['fixTime'])
        df = df.sort_values('timestamp')
        
        # Calculate speed and bearing
        df['speed_kmh'] = df['speed'] * 1.852  # Convert knots to km/h
        df['elevation'] = df.get('altitude', 0)
        
        # Calculate distance between points
        distances = []
        bearings = []
        
        for i in range(len(df)):
            if i == 0:
                distances.append(0)
                bearings.append(0)
            else:
                prev_point = (df.iloc[i-1]['latitude'], df.iloc[i-1]['longitude'])
                curr_point = (df.iloc[i]['latitude'], df.iloc[i]['longitude'])
                
                # Calculate distance
                dist = geodesic(prev_point, curr_point).meters
                distances.append(dist)
                
                # Calculate bearing
                lat1, lon1 = math.radians(prev_point[0]), math.radians(prev_point[1])
                lat2, lon2 = math.radians(curr_point[0]), math.radians(curr_point[1])
                
                dlon = lon2 - lon1
                bearing = math.atan2(
                    math.sin(dlon) * math.cos(lat2),
                    math.cos(lat1) * math.sin(lat2) - math.sin(lat1) * math.cos(lat2) * math.cos(dlon)
                )
                bearings.append(math.degrees(bearing))
        
        df['distance'] = distances
        df['bearing'] = bearings
        
        return df
    
    def create_pydeck_frame(self, df, current_index, frame_num):
        """Create a single frame using Pydeck"""
        # Get current position
        current_row = df.iloc[current_index]
        
        # Get trail data (previous points)
        start_idx = max(0, current_index - self.trail_length)
        trail_data = df.iloc[start_idx:current_index + 1].copy()
        
        # Create color gradient for trail (fade effect)
        trail_colors = []
        for i, _ in enumerate(trail_data.iterrows()):
            alpha = (i + 1) / len(trail_data) * 255
            trail_colors.append([255, 100, 100, int(alpha)])
        
        trail_data['color'] = trail_colors
        
        # Create the deck
        view_state = pdk.ViewState(
            longitude=current_row['longitude'],
            latitude=current_row['latitude'],
            zoom=16,
            pitch=60,
            bearing=current_row['bearing']
        )
        
        # Path layer (trail)
        path_layer = pdk.Layer(
            "PathLayer",
            trail_data,
            get_path=lambda x: [[x['longitude'], x['latitude'], x['elevation']]],
            get_color=[255, 100, 100, 200],
            width_min_pixels=3,
            width_max_pixels=8,
        )
        
        # Scatterplot layer (current position)
        current_point = pd.DataFrame([{
            'longitude': current_row['longitude'],
            'latitude': current_row['latitude'],
            'elevation': current_row['elevation'] + 10,
            'speed': current_row['speed_kmh']
        }])
        
        scatter_layer = pdk.Layer(
            "ScatterplotLayer",
            current_point,
            get_position=['longitude', 'latitude', 'elevation'],
            get_radius=15,
            get_color=[255, 255, 0, 255],
            pickable=True
        )
        
        # Create deck
        deck = pdk.Deck(
            layers=[path_layer, scatter_layer],
            initial_view_state=view_state,
            map_style='mapbox://styles/mapbox/satellite-v9'
        )
        
        # Save frame
        frame_path = os.path.join(self.frames_folder, f"frame_{frame_num:06d}.png")
        deck.to_html(frame_path.replace('.png', '.html'))
        
        return frame_path
    
    def create_plotly_frame(self, df, current_index, frame_num):
        """Create a single frame using Plotly for 3D visualization"""
        # Get current position
        current_row = df.iloc[current_index]
        
        # Get trail data
        start_idx = max(0, current_index - self.trail_length)
        trail_data = df.iloc[start_idx:current_index + 1]
        
        # Create 3D scatter plot
        fig = go.Figure()
        
        # Add trail
        fig.add_trace(go.Scatter3d(
            x=trail_data['longitude'],
            y=trail_data['latitude'],
            z=trail_data['elevation'],
            mode='lines+markers',
            line=dict(
                color=trail_data.index,
                colorscale='Plasma',
                width=8
            ),
            marker=dict(
                size=3,
                opacity=0.8
            ),
            name='Trail'
        ))
        
        # Add current position
        fig.add_trace(go.Scatter3d(
            x=[current_row['longitude']],
            y=[current_row['latitude']],
            z=[current_row['elevation'] + 50],
            mode='markers',
            marker=dict(
                size=15,
                color='red',
                symbol='diamond'
            ),
            name='Current Position'
        ))
        
        # Add speed information as text
        speed_text = f"Speed: {current_row['speed_kmh']:.1f} km/h<br>"
        speed_text += f"Time: {current_row['timestamp'].strftime('%H:%M:%S')}<br>"
        speed_text += f"Altitude: {current_row['elevation']:.0f} m"
        
        # Update layout
        fig.update_layout(
            title=f"3D GPS Track Animation - Frame {frame_num}",
            scene=dict(
                xaxis_title='Longitude',
                yaxis_title='Latitude',
                zaxis_title='Elevation (m)',
                camera=dict(
                    eye=dict(x=1.5, y=1.5, z=1.5)
                ),
                aspectmode='cube'
            ),
            annotations=[
                dict(
                    text=speed_text,
                    x=0.02,
                    y=0.98,
                    xref='paper',
                    yref='paper',
                    showarrow=False,
                    bgcolor='rgba(255,255,255,0.8)',
                    bordercolor='black',
                    borderwidth=1
                )
            ],
            width=1920,
            height=1080
        )
        
        # Save frame
        frame_path = os.path.join(self.frames_folder, f"frame_{frame_num:06d}.png")
        fig.write_image(frame_path, engine="kaleido")
        
        return frame_path
    
    def create_advanced_plotly_frame(self, df, current_index, frame_num):
        """Create Relive-style progressive animation frame"""
        # Progressive track: show all points from start to current position
        current_track = df.iloc[:current_index + 1]
        current_row = df.iloc[current_index]
        
        # Create subplot with multiple views
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=[
                f'3D Track Progress - Frame {frame_num + 1}', 
                'Route Overview', 
                'Speed Over Time', 
                'Elevation Profile'
            ],
            specs=[[{'type': 'scene'}, {'type': 'scatter'}],
                   [{'type': 'scatter'}, {'type': 'scatter'}]],
            vertical_spacing=0.1,
            horizontal_spacing=0.1
        )
        
        # 1. 3D Progressive Track View
        if len(current_track) > 1:
            # Show completed track in blue
            fig.add_trace(
                go.Scatter3d(
                    x=current_track['longitude'],
                    y=current_track['latitude'], 
                    z=current_track['elevation'],
                    mode='lines+markers',
                    line=dict(color='blue', width=5),
                    marker=dict(
                        size=3, 
                        color=current_track['speed_kmh'], 
                        colorscale='Viridis',
                        showscale=True,
                        colorbar=dict(title="Speed (km/h)", x=0.45)
                    ),
                    name='Completed Track',
                    showlegend=False
                ),
                row=1, col=1
            )
        
        # Current vehicle position (moving marker)
        fig.add_trace(
            go.Scatter3d(
                x=[current_row['longitude']],
                y=[current_row['latitude']],
                z=[current_row['elevation'] + 15],
                mode='markers',
                marker=dict(
                    size=15, 
                    color='red', 
                    symbol='diamond',
                    line=dict(color='white', width=2)
                ),
                name=f'Vehicle - {current_row["timestamp"].strftime("%H:%M:%S")}',
                showlegend=False
            ),
            row=1, col=1
        )
        
        # 2. Top View - Route Overview with full track
        # Show full route in light gray
        fig.add_trace(
            go.Scatter(
                x=df['longitude'],
                y=df['latitude'],
                mode='lines',
                line=dict(color='lightgray', width=2, dash='dot'),
                name='Full Route',
                showlegend=False
            ),
            row=1, col=2
        )
        
        # Show completed track in color
        if len(current_track) > 1:
            fig.add_trace(
                go.Scatter(
                    x=current_track['longitude'],
                    y=current_track['latitude'],
                    mode='lines+markers',
                    line=dict(color='blue', width=4),
                    marker=dict(
                        size=4, 
                        color=current_track['speed_kmh'], 
                        colorscale='Viridis'
                    ),
                    name='Progress',
                    showlegend=False
                ),
                row=1, col=2
            )
        
        # Current position on map
        fig.add_trace(
            go.Scatter(
                x=[current_row['longitude']],
                y=[current_row['latitude']],
                mode='markers',
                marker=dict(
                    size=12, 
                    color='red',
                    symbol='circle',
                    line=dict(color='white', width=3)
                ),
                name='Current Position',
                showlegend=False
            ),
            row=1, col=2
        )
        
        # 3. Speed Profile Over Time
        if len(current_track) > 1:
            time_points = list(range(len(current_track)))
            fig.add_trace(
                go.Scatter(
                    x=time_points,
                    y=current_track['speed_kmh'],
                    mode='lines+markers',
                    line=dict(color='green', width=3),
                    marker=dict(size=4),
                    fill='tonexty',
                    name='Speed',
                    showlegend=False
                ),
                row=2, col=1
            )
            
            # Current speed marker
            fig.add_trace(
                go.Scatter(
                    x=[current_index],
                    y=[current_row['speed_kmh']],
                    mode='markers',
                    marker=dict(size=10, color='red'),
                    name='Current Speed',
                    showlegend=False
                ),
                row=2, col=1
            )
        
        # 4. Elevation Profile Over Time
        if len(current_track) > 1:
            time_points = list(range(len(current_track)))
            fig.add_trace(
                go.Scatter(
                    x=time_points,
                    y=current_track['elevation'],
                    mode='lines+markers',
                    line=dict(color='brown', width=3),
                    marker=dict(size=4),
                    fill='tonexty',
                    name='Elevation',
                    showlegend=False
                ),
                row=2, col=2
            )
            
            # Current elevation marker
            fig.add_trace(
                go.Scatter(
                    x=[current_index],
                    y=[current_row['elevation']],
                    mode='markers',
                    marker=dict(size=10, color='red'),
                    name='Current Elevation',
                    showlegend=False
                ),
                row=2, col=2
            )
        
        # Enhanced layout with better styling
        fig.update_layout(
            title=dict(
                text=f'GPS Track Animation - {current_row["timestamp"].strftime("%Y-%m-%d %H:%M:%S")}<br>' +
                     f'Progress: {current_index + 1}/{len(df)} points ({((current_index + 1)/len(df)*100):.1f}%)',
                x=0.5,
                font=dict(size=16)
            ),
            showlegend=False,
            width=1920,
            height=1080,
            paper_bgcolor='white',
            plot_bgcolor='white'
        )
        
        # Update 3D scene
        fig.update_scenes(
            camera=dict(
                eye=dict(x=1.5, y=1.5, z=1.2),
                center=dict(x=0, y=0, z=0),
                up=dict(x=0, y=0, z=1)
            ),
            aspectmode='cube'
        )
        
        # Update axes labels
        fig.update_xaxes(title_text="Longitude", row=1, col=2)
        fig.update_yaxes(title_text="Latitude", row=1, col=2)
        fig.update_xaxes(title_text="Time Points", row=2, col=1)
        fig.update_yaxes(title_text="Speed (km/h)", row=2, col=1)
        fig.update_xaxes(title_text="Time Points", row=2, col=2)
        fig.update_yaxes(title_text="Elevation (m)", row=2, col=2)
        
        # Save frame
        frame_path = os.path.join(self.frames_folder, f"advanced_3d_frame_{frame_num:04d}.png")
        
        try:
            fig.write_image(frame_path, engine="kaleido", width=1920, height=1080)
            return frame_path
        except Exception as e:
            print(f"Error creating frame {frame_num}: {e}")
            return None
    def generate_frames(self, positions_file, style='advanced', progress_callback=None):
        """Generate Relive-style progressive animation frames"""
        print("Loading GPS data...")
        df = self.load_gps_data(positions_file)
        
        if len(df) < 2:
            print("Not enough GPS points for animation")
            return []
        
        # Animation settings for smooth progression
        min_frames = 60  # Minimum frames for very short trips
        max_frames = 300  # Maximum frames to keep file size reasonable
        
        # Calculate optimal frame count based on trip length
        total_frames = min(max_frames, max(min_frames, len(df) * 2))
        
        # Calculate step size for smooth progression
        step_size = len(df) / total_frames
        
        frame_paths = []
        
        print(f"Generating {total_frames} frames for Relive-style animation...")
        print(f"Processing {len(df)} GPS points with step size {step_size:.2f}")
        
        for frame_num in range(total_frames):
            # Calculate which GPS point to show up to
            current_index = min(int(frame_num * step_size), len(df) - 1)
            
            # Ensure we always progress forward
            if current_index == 0:
                current_index = 1
            
            try:
                if style == 'pydeck':
                    frame_path = self.create_pydeck_frame(df, current_index, frame_num)
                elif style == 'plotly':
                    frame_path = self.create_plotly_frame(df, current_index, frame_num)
                else:  # advanced
                    frame_path = self.create_advanced_plotly_frame(df, current_index, frame_num)
                
                if frame_path:
                    frame_paths.append(frame_path)
                
                # Update progress
                if progress_callback:
                    progress = ((frame_num + 1) / total_frames) * 100
                    progress_callback(
                        progress, 
                        f"Creating animation frame {frame_num + 1}/{total_frames} (GPS point {current_index + 1}/{len(df)})"
                    )
                
                # Progress feedback
                if (frame_num + 1) % 20 == 0:
                    print(f"Generated {frame_num + 1}/{total_frames} frames ({progress:.1f}%)")
                    
            except Exception as e:
                print(f"Error generating frame {frame_num}: {e}")
                continue
        
        print(f"Successfully generated {len(frame_paths)} animation frames")
        return frame_paths
    
    def create_video(self, frame_paths, output_video_path, progress_callback=None):
        """Create video from frames using MoviePy with optimized settings"""
        print("Creating Relive-style animation video...")
        
        if not frame_paths:
            print("No frames to create video from")
            return False
        
        try:
            # Filter out None paths
            valid_frames = [f for f in frame_paths if f and os.path.exists(f)]
            
            if not valid_frames:
                print("No valid frames found")
                return False
            
            print(f"Creating video from {len(valid_frames)} frames at {self.fps} FPS...")
            
            # Create video clip from images with optimal settings
            clip = ImageSequenceClip(valid_frames, fps=self.fps)
            
            # Add smooth fade effects for professional look
            clip = clip.fadein(0.5).fadeout(0.5)
            
            # Update progress
            if progress_callback:
                progress_callback(50, "Encoding video with optimized settings...")
            
            # Write video file with high quality settings
            clip.write_videofile(
                output_video_path,
                codec='libx264',
                audio=False,
                temp_audiofile=None,
                remove_temp=True,
                verbose=False,
                logger=None,
                bitrate="8000k",  # High quality bitrate
                ffmpeg_params=[
                    "-preset", "medium",  # Balance between speed and compression
                    "-crf", "18",  # High quality (lower = better quality)
                    "-pix_fmt", "yuv420p"  # Better compatibility
                ]
            )
            
            if progress_callback:
                progress_callback(100, f"Video successfully created: {os.path.basename(output_video_path)}")
            
            print(f"✅ Relive-style animation video saved to: {output_video_path}")
            print(f"📊 Video info: {len(valid_frames)} frames, {self.fps} FPS, {clip.duration:.1f}s duration")
            
            # Clean up clip
            clip.close()
            
            return True
            
        except Exception as e:
            print(f"Error creating video: {e}")
            return False
    
    def cleanup_frames(self):
        """Clean up temporary frame files"""
        import shutil
        if os.path.exists(self.frames_folder):
            shutil.rmtree(self.frames_folder)
        os.makedirs(self.frames_folder, exist_ok=True)
    
    def generate_3d_animation(self, positions_file, output_video_path, 
                            style='advanced', cleanup=True, progress_callback=None):
        """
        Main method to generate 3D animation
        
        Args:
            positions_file: Path to JSON file with GPS positions
            output_video_path: Path for output video
            style: 'pydeck', 'plotly', or 'advanced'
            cleanup: Whether to clean up temporary files
            progress_callback: Callback function for progress updates
        """
        try:
            # Generate frames
            frame_paths = self.generate_frames(positions_file, style, progress_callback)
            
            if not frame_paths:
                raise Exception("No frames generated")
            
            # Create video
            success = self.create_video(frame_paths, output_video_path, progress_callback)
            
            if cleanup:
                self.cleanup_frames()
            
            return success
            
        except Exception as e:
            print(f"Error generating 3D animation: {e}")
            if progress_callback:
                progress_callback(-1, f"Error: {e}")
            return False

def generate_advanced_3d_video(positions_file, output_folder, filename_prefix="advanced_3d", 
                             style='advanced', progress_callback=None):
    """
    Convenience function to generate advanced 3D video
    """
    generator = Advanced3DGenerator(output_folder)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_video_path = os.path.join(output_folder, f"{filename_prefix}_{timestamp}.mp4")
    
    success = generator.generate_3d_animation(
        positions_file, 
        output_video_path, 
        style=style,
        progress_callback=progress_callback
    )
    
    return output_video_path if success else None

# Test function
if __name__ == "__main__":
    # Test the advanced 3D generator
    test_positions = "test_positions.json"
    output_dir = "test_output"
    
    def test_progress(progress, message):
        print(f"Progress: {progress:.1f}% - {message}")
    
    video_path = generate_advanced_3d_video(
        test_positions, 
        output_dir, 
        style='advanced',
        progress_callback=test_progress
    )
    
    if video_path:
        print(f"Test video created: {video_path}")
    else:
        print("Test failed")