"""
Utility script to project point clouds or meshes onto the Z plane (XY plane)
and compare the overlap between fish point cloud projection and template point cloud projection.

The overlap is calculated as: intersection_area / fish_projection_area * 100%

Classification:
- High overlap: >= 90%
- Medium overlap: 70% - 90%
- Low overlap: < 70%
"""

import numpy as np
import trimesh
import open3d as o3d
from pathlib import Path
from typing import Union, Tuple, Optional
from scipy.spatial import ConvexHull

# Try importing shapely, but provide fallback if not available
try:
    from shapely.geometry import Polygon, Point
    from shapely.ops import unary_union
    SHAPELY_AVAILABLE = True
except ImportError:
    SHAPELY_AVAILABLE = False
    print("Warning: shapely not available. Using fallback method for polygon operations.")


def load_mesh_or_pointcloud(file_path: Union[str, Path]):
    """
    Load a mesh or point cloud from file.
    
    Args:
        file_path: Path to PLY, OBJ, or STL file
        
    Returns:
        numpy array of shape (N, 3) containing vertices/points
    """
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"File not found: {file_path}")
    
    try:
        # Try loading as mesh first
        mesh = trimesh.load(str(file_path))
        if hasattr(mesh, 'vertices'):
            return np.asarray(mesh.vertices)
        elif hasattr(mesh, 'points'):
            return np.asarray(mesh.points)
        else:
            raise ValueError(f"Unable to extract points from {file_path}")
    except Exception as e:
        # Fallback to open3d
        try:
            mesh_o3d = o3d.io.read_triangle_mesh(str(file_path))
            if len(mesh_o3d.vertices) > 0:
                return np.asarray(mesh_o3d.vertices)
            else:
                # Try as point cloud
                pcd = o3d.io.read_point_cloud(str(file_path))
                if len(pcd.points) > 0:
                    return np.asarray(pcd.points)
                else:
                    raise ValueError(f"No points found in {file_path}")
        except Exception as e2:
            raise ValueError(f"Failed to load {file_path}: {e}, {e2}")


def remove_tail_portion(points: np.ndarray, tail_ratio: float = 0.2) -> np.ndarray:
    """
    Remove the tail portion (negative X direction) from points.
    Removes points in the tail_ratio (e.g., 0.2 = 20%) of the length in negative X direction.
    
    Args:
        points: Nx3 array of 3D points
        tail_ratio: Ratio of tail to remove (default: 0.2 for 20%)
        
    Returns:
        Filtered Nx3 array of points with tail removed
    """
    if points.shape[1] != 3:
        raise ValueError(f"Expected Nx3 array, got shape {points.shape}")
    
    if len(points) == 0:
        return points
    
    # Calculate X-axis range
    x_vals = points[:, 0]
    min_x = np.min(x_vals)
    max_x = np.max(x_vals)
    length = max_x - min_x
    
    # Calculate threshold: remove points in the tail_ratio portion from the negative X side
    # Keep points where X >= (min_x + tail_ratio * length)
    threshold = min_x + tail_ratio * length
    
    # Filter points
    mask = x_vals >= threshold
    filtered_points = points[mask]
    
    return filtered_points


def project_to_xy_plane(points: np.ndarray) -> np.ndarray:
    """
    Project 3D points onto the XY plane (Z=0).
    
    Args:
        points: Nx3 array of 3D points
        
    Returns:
        Nx2 array of 2D points (X, Y coordinates)
    """
    if points.shape[1] != 3:
        raise ValueError(f"Expected Nx3 array, got shape {points.shape}")
    
    return points[:, :2]


def points_to_polygon(points_2d: np.ndarray, method: str = 'convex_hull'):
    """
    Convert 2D points to a polygon (Shapely Polygon if available, otherwise numpy array).
    
    Args:
        points_2d: Nx2 array of 2D points
        method: Method to use ('convex_hull')
        
    Returns:
        Shapely Polygon object if shapely is available, otherwise numpy array of hull points
    """
    if len(points_2d) < 3:
        raise ValueError("Need at least 3 points to create a polygon")
    
    if method == 'convex_hull':
        # Use convex hull
        try:
            hull = ConvexHull(points_2d)
            hull_points = points_2d[hull.vertices]
            
            if SHAPELY_AVAILABLE:
                polygon = Polygon(hull_points)
                if not polygon.is_valid:
                    polygon = polygon.buffer(0)  # Fix invalid polygon
                return polygon
            else:
                # Return hull points as numpy array (closed polygon)
                return np.vstack([hull_points, hull_points[0]])
        except Exception as e:
            # Fallback: create a simple bounding box
            min_x, min_y = points_2d.min(axis=0)
            max_x, max_y = points_2d.max(axis=0)
            bbox_points = np.array([
                [min_x, min_y],
                [max_x, min_y],
                [max_x, max_y],
                [min_x, max_y],
                [min_x, min_y]  # Close the polygon
            ])
            
            if SHAPELY_AVAILABLE:
                return Polygon(bbox_points[:-1])  # Exclude duplicate closing point
            else:
                return bbox_points
    else:
        raise ValueError(f"Unknown method: {method}")


def polygon_area_numpy(polygon_points: np.ndarray) -> float:
    """
    Calculate polygon area using the shoelace formula.
    
    Args:
        polygon_points: Nx2 array of polygon vertices (closed)
        
    Returns:
        Area in square units
    """
    # Remove duplicate closing point if present
    if len(polygon_points) > 1 and np.allclose(polygon_points[0], polygon_points[-1]):
        polygon_points = polygon_points[:-1]
    
    x = polygon_points[:, 0]
    y = polygon_points[:, 1]
    # Shoelace formula
    area = 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))
    return area


def polygon_intersection_area_numpy(poly1: np.ndarray, poly2: np.ndarray) -> float:
    """
    Calculate intersection area of two polygons using a grid-based method.
    This is a fallback when shapely is not available.
    
    Args:
        poly1: Nx2 array of first polygon vertices
        poly2: Mx2 array of second polygon vertices
        
    Returns:
        Intersection area
    """
    try:
        from matplotlib.path import Path as MplPath
    except ImportError:
        # If matplotlib is not available, use a simple approximation
        # Calculate bounding box intersection as approximation
        min1 = poly1.min(axis=0)
        max1 = poly1.max(axis=0)
        min2 = poly2.min(axis=0)
        max2 = poly2.max(axis=0)
        
        # Intersection of bounding boxes
        inter_min = np.maximum(min1, min2)
        inter_max = np.minimum(max1, max2)
        
        if np.any(inter_min >= inter_max):
            return 0.0
        
        return np.prod(inter_max - inter_min)
    
    # Find bounding box of both polygons
    all_points = np.vstack([poly1, poly2])
    min_x, min_y = all_points.min(axis=0)
    max_x, max_y = all_points.max(axis=0)
    
    # Create a grid
    grid_resolution = 200  # Adjust based on desired accuracy
    x_grid = np.linspace(min_x, max_x, grid_resolution)
    y_grid = np.linspace(min_y, max_y, grid_resolution)
    cell_area = ((max_x - min_x) / grid_resolution) * ((max_y - min_y) / grid_resolution)
    
    # Check which grid points are inside both polygons
    path1 = MplPath(poly1)
    path2 = MplPath(poly2)
    
    intersection_count = 0
    for x in x_grid:
        for y in y_grid:
            if path1.contains_point((x, y)) and path2.contains_point((x, y)):
                intersection_count += 1
    
    return intersection_count * cell_area


def calculate_overlap(fish_points: np.ndarray, 
                     template_points: np.ndarray,
                     method: str = 'convex_hull',
                     remove_tail: bool = True,
                     tail_ratio: float = 0.2) -> Tuple[float, float, float, str]:
    """
    Calculate overlap between fish point cloud projection and template point cloud projection.
    
    Args:
        fish_points: Nx3 array of fish point cloud
        template_points: Mx3 array of template point cloud
        method: Method for polygon creation ('convex_hull')
        remove_tail: Whether to remove tail portion (20% in negative X direction) before calculation
        tail_ratio: Ratio of tail to remove if remove_tail is True (default: 0.2 for 20%)
        
    Returns:
        tuple: (overlap_percent, intersection_area, fish_area, template_area, classification)
        - overlap_percent: Intersection over Fish percentage (0-100)
        - intersection_area: Area of intersection in mm²
        - fish_area: Area of fish projection in mm²
        - template_area: Area of template projection in mm²
        - classification: 'high', 'medium', or 'low'
    """
    # Remove tail portion if requested
    if remove_tail:
        fish_points = remove_tail_portion(fish_points, tail_ratio=tail_ratio)
        template_points = remove_tail_portion(template_points, tail_ratio=tail_ratio)
    
    # Project to XY plane
    fish_2d = project_to_xy_plane(fish_points)
    template_2d = project_to_xy_plane(template_points)
    
    # Convert to polygons
    fish_polygon = points_to_polygon(fish_2d, method=method)
    template_polygon = points_to_polygon(template_2d, method=method)
    
    # Calculate areas
    if SHAPELY_AVAILABLE:
        fish_area = fish_polygon.area
        template_area = template_polygon.area
        
        # Calculate intersection
        try:
            intersection = fish_polygon.intersection(template_polygon)
            if intersection.is_empty:
                intersection_area = 0.0
            else:
                # Handle MultiPolygon or other geometry types
                if hasattr(intersection, 'area'):
                    intersection_area = intersection.area
                else:
                    # If it's a collection, sum all areas
                    intersection_area = sum(p.area for p in intersection.geoms) if hasattr(intersection, 'geoms') else 0.0
        except Exception as e:
            print(f"Warning: Error calculating intersection: {e}")
            intersection_area = 0.0
    else:
        # Fallback: use numpy-based calculations
        fish_area = polygon_area_numpy(fish_polygon)
        template_area = polygon_area_numpy(template_polygon)
        intersection_area = polygon_intersection_area_numpy(fish_polygon, template_polygon)
    
    # Calculate overlap percentage using Intersection over Fish
    if fish_area > 0:
        overlap_percent = (intersection_area / fish_area) * 100.0
    else:
        overlap_percent = 0.0
    
    # Classify overlap
    if overlap_percent >= 90.0:
        classification = 'high'
    elif overlap_percent >= 70.0:
        classification = 'medium'
    else:
        classification = 'low'
    
    return overlap_percent, intersection_area, fish_area, template_area, classification


def compare_overlap(fish_file: Union[str, Path],
                   template_file: Union[str, Path],
                   method: str = 'convex_hull',
                   verbose: bool = True,
                   remove_tail: bool = True,
                   tail_ratio: float = 0.2) -> dict:
    """
    Compare overlap between fish and template point clouds/meshes.
    
    Args:
        fish_file: Path to fish point cloud/mesh file
        template_file: Path to template point cloud/mesh file
        method: Method for polygon creation ('convex_hull')
        verbose: Print detailed information
        remove_tail: Whether to remove tail portion (20% in negative X direction) before calculation
        tail_ratio: Ratio of tail to remove if remove_tail is True (default: 0.2 for 20%)
        
    Returns:
        Dictionary with overlap information:
        {
            'overlap_percent': float,  # Intersection over Fish percentage
            'intersection_area_mm2': float,
            'fish_area_mm2': float,
            'template_area_mm2': float,
            'classification': str ('high', 'medium', or 'low'),
            'fish_file': str,
            'template_file': str
        }
    """
    # Load point clouds/meshes
    fish_points = load_mesh_or_pointcloud(fish_file)
    template_points = load_mesh_or_pointcloud(template_file)
    
    if verbose:
        print(f"Fish point cloud: {len(fish_points)} points")
        print(f"Template point cloud: {len(template_points)} points")
        if remove_tail:
            print(f"Removing tail portion ({tail_ratio*100:.0f}% in negative X direction)")
    
    # Calculate overlap
    overlap_percent, intersection_area, fish_area, template_area, classification = calculate_overlap(
        fish_points, template_points, method=method, remove_tail=remove_tail, tail_ratio=tail_ratio
    )
    
    result = {
        'overlap_percent': overlap_percent,
        'intersection_area_mm2': intersection_area,
        'fish_area_mm2': fish_area,
        'template_area_mm2': template_area,
        'classification': classification,
        'fish_file': str(fish_file),
        'template_file': str(template_file)
    }
    
    if verbose:
        print(f"\n{'='*60}")
        print("Overlap Analysis")
        print(f"{'='*60}")
        print(f"Fish projection area: {fish_area:.2f} mm²")
        print(f"Template projection area: {template_area:.2f} mm²")
        print(f"Intersection area: {intersection_area:.2f} mm²")
        print(f"Overlap (Intersection/Fish): {overlap_percent:.2f}%")
        print(f"Classification: {classification.upper()} overlap")
        if classification == 'low':
            print(f"⚠ Warning: Low overlap detected. Template may need rescaling.")
        print(f"{'='*60}")
    
    return result


def main():
    """Command-line interface for overlap comparison."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Compare overlap between fish and template point cloud projections"
    )
    parser.add_argument(
        "--fish", "-f",
        type=str,
        required=True,
        help="Path to fish point cloud/mesh file"
    )
    parser.add_argument(
        "--template", "-t",
        type=str,
        required=True,
        help="Path to template point cloud/mesh file"
    )
    parser.add_argument(
        "--method",
        type=str,
        default="convex_hull",
        choices=["convex_hull"],
        help="Method for polygon creation (default: convex_hull)"
    )
    parser.add_argument(
        "--output", "-o",
        type=str,
        default=None,
        help="Output JSON file to save results"
    )
    parser.add_argument(
        "--quiet", "-q",
        action="store_true",
        help="Suppress verbose output"
    )
    parser.add_argument(
        "--remove-tail",
        action="store_true",
        default=True,
        help="Remove tail portion (20% in negative X direction) before calculation (default: True)"
    )
    parser.add_argument(
        "--no-remove-tail",
        action="store_false",
        dest="remove_tail",
        help="Do not remove tail portion"
    )
    parser.add_argument(
        "--tail-ratio",
        type=float,
        default=0.2,
        help="Ratio of tail to remove (default: 0.2 for 20%%)"
    )
    
    args = parser.parse_args()
    
    # Compare overlap
    result = compare_overlap(
        args.fish,
        args.template,
        method=args.method,
        verbose=not args.quiet,
        remove_tail=args.remove_tail,
        tail_ratio=args.tail_ratio
    )
    
    # Save to JSON if requested
    if args.output:
        import json
        with open(args.output, 'w') as f:
            json.dump(result, f, indent=2)
        print(f"\nResults saved to: {args.output}")
    
    return result


if __name__ == "__main__":
    main()