"""
isea3hcompact.py - ISEA3H Cell Compaction Utilities

This module provides functions and command-line interfaces for compacting and expanding ISEA3H cells.
It supports flexible input and output formats, including file paths (GeoJSON, Shapefile, CSV, Parquet),
GeoDataFrames, lists of cell IDs, and GeoJSON dictionaries. Outputs can be written to various formats or
returned as Python objects. The main functions are:

- isea3hcompact: Compact a set of ISEA3H cells to their minimal covering set.
- isea3hexpand: Expand (uncompact) a set of ISEA3H cells to a target resolution.
- isea3hcompact_cli: Command-line interface for compaction.
- isea3hexpand_cli: Command-line interface for expansion.

Dependencies: geopandas, pandas, shapely, vgrid.dggs.eaggr, vgrid DGGS.
"""
import os
import argparse
import platform
import geopandas as gpd
from collections import defaultdict

if platform.system() == "Windows":
    from vgrid.dggs.eaggr.eaggr import Eaggr
    from vgrid.dggs.eaggr.shapes.dggs_cell import DggsCell
    from vgrid.dggs.eaggr.enums.model import Model
    from vgrid.utils.constants import ISEA3H_ACCURACY_RES_DICT
    isea3h_dggs = Eaggr(Model.ISEA3H)

from vgrid.conversion.dggs2geo.isea3h2geo import isea3h2geo
from vgrid.utils.geometry import geodesic_dggs_to_geoseries, isea3h_cell_to_polygon
from vgrid.utils.io import process_input_data_compact, convert_to_output_format
from pyproj import Geod

geod = Geod(ellps="WGS84")

# --- ISEA3H Compaction/Expansion Logic ---
def get_isea3h_cell_children(isea3h_cell, resolution):
    """Recursively expands a DGGS cell until all children reach the desired resolution."""
    if platform.system() != "Windows":
        raise RuntimeError("ISEA3H compaction is only supported on Windows.")
    
    isea3h2point = isea3h_dggs.convert_dggs_cell_to_point(isea3h_cell)
    cell_accuracy = isea3h2point._accuracy
    cell_resolution = ISEA3H_ACCURACY_RES_DICT.get(cell_accuracy)

    if cell_resolution >= resolution:
        return [isea3h_cell]  # Base case: return the cell if it meets/exceeds resolution

    expanded_cells = []
    children = isea3h_dggs.get_dggs_cell_children(isea3h_cell)

    for child in children:
        expanded_cells.extend(get_isea3h_cell_children(child, resolution))

    return expanded_cells

def get_isea3h_resolution(isea3h_id):
    """Get the resolution of an ISEA3H cell ID."""
    try:
        isea3h_cell = DggsCell(isea3h_id)
        cell_polygon = isea3h_cell_to_polygon(isea3h_cell)
        cell_perimeter = abs(geod.geometry_area_perimeter(cell_polygon)[1])

        isea3h2point = isea3h_dggs.convert_dggs_cell_to_point(isea3h_cell)
        cell_accuracy = isea3h2point._accuracy

        avg_edge_len = cell_perimeter / 6
        cell_resolution = ISEA3H_ACCURACY_RES_DICT.get(cell_accuracy)

        if cell_resolution == 0:  # icosahedron faces at resolution = 0
            avg_edge_len = cell_perimeter / 3

        if cell_accuracy == 0.0:
            if round(avg_edge_len, 2) == 0.06:
                cell_resolution = 33
            elif round(avg_edge_len, 2) == 0.03:
                cell_resolution = 34
            elif round(avg_edge_len, 2) == 0.02:
                cell_resolution = 35
            elif round(avg_edge_len, 2) == 0.01:
                cell_resolution = 36
            elif round(avg_edge_len, 3) == 0.007:
                cell_resolution = 37
            elif round(avg_edge_len, 3) == 0.004:
                cell_resolution = 38
            elif round(avg_edge_len, 3) == 0.002:
                cell_resolution = 39
            elif round(avg_edge_len, 3) <= 0.001:
                cell_resolution = 40

        return cell_resolution
    except Exception as e:
        raise ValueError(f"Invalid cell ID <{isea3h_id}> : {e}")

def isea3h_compact(isea3h_ids):
    """Compact a list of ISEA3H cell IDs to their minimal covering set."""
    if platform.system() != "Windows":
        raise RuntimeError("ISEA3H compaction is only supported on Windows.")
    
    isea3h_ids = set(isea3h_ids)  # Remove duplicates
    cell_cache = {cell_id: DggsCell(cell_id) for cell_id in isea3h_ids}

    while True:
        grouped_by_parent = defaultdict(set)

        # Group cells by *all* their parents
        for cell_id in isea3h_ids:
            cell = cell_cache[cell_id]
            try:
                parents = isea3h_dggs.get_dggs_cell_parents(cell)
            except Exception as e:
                print(f"Error getting parents for {cell_id}: {e}")
                continue

            for parent in parents:
                parent_id = parent.get_cell_id()
                grouped_by_parent[parent_id].add(cell_id)

        new_isea3h_ids = set(isea3h_ids)
        changed = False

        for parent_id, children_ids in grouped_by_parent.items():
            parent_cell = DggsCell(parent_id)
            try:
                expected_children = set(
                    child.get_cell_id()
                    for child in isea3h_dggs.get_dggs_cell_children(parent_cell)
                )
            except Exception as e:
                print(f"Error getting children for parent {parent_id}: {e}")
                continue

            # Check for full match: only then compact
            if children_ids == expected_children:
                new_isea3h_ids.difference_update(children_ids)
                new_isea3h_ids.add(parent_id)
                cell_cache[parent_id] = parent_cell
                changed = True
            else:
                # Keep original children if they don't fully match expected subcells
                new_isea3h_ids.update(children_ids)

        if not changed:
            break  # Fully compacted

        isea3h_ids = new_isea3h_ids

    return sorted(isea3h_ids)

def isea3hcompact(
    input_data,
    isea3h_id=None,
    output_format=None,
):
    """Compact ISEA3H cells from input data."""
    if platform.system() != "Windows":
        raise RuntimeError("ISEA3H compaction is only supported on Windows.")
    
    if not isea3h_id:
        isea3h_id = "isea3h"
    
    gdf = process_input_data_compact(input_data, isea3h_id)
    isea3h_ids = gdf[isea3h_id].drop_duplicates().tolist()
    
    if not isea3h_ids:
        print(f"No ISEA3H IDs found in <{isea3h_id}> field.")
        return
    
    try:
        isea3h_ids_compact = isea3h_compact(isea3h_ids)
    except Exception:
        raise Exception("Compact cells failed. Please check your ISEA3H ID field.")
    
    if not isea3h_ids_compact:
        return None
    
    rows = []
    for isea3h_id_compact in isea3h_ids_compact:
        try:
            cell_polygon = isea3h2geo(isea3h_id_compact)
            cell_resolution = get_isea3h_resolution(isea3h_id_compact)
            num_edges = 6  # ISEA3H cells are hexagonal
            row = geodesic_dggs_to_geoseries(
                "isea3h", isea3h_id_compact, cell_resolution, cell_polygon, num_edges
            )
            rows.append(row)
        except Exception:
            continue
    
    out_gdf = gpd.GeoDataFrame(rows, geometry="geometry",crs="EPSG:4326")
    
    file_formats = ["csv", "geojson", "shapefile", "gpkg", "parquet", "geoparquet"]
    output_name = None
    if output_format in file_formats:
        ext_map = {
            "csv": ".csv",
            "geojson": ".geojson",
            "shapefile": ".shp",
            "gpkg": ".gpkg",
            "parquet": ".parquet",
            "geoparquet": ".parquet",
        }
        ext = ext_map.get(output_format, "")
        if isinstance(input_data, str):
            base = os.path.splitext(os.path.basename(input_data))[0]
            output_name = f"{base}_isea3h_compacted{ext}"
        else:
            output_name = f"isea3h_compacted{ext}"
    
    return convert_to_output_format(out_gdf, output_format, output_name)

def isea3hcompact_cli():
    """Command-line interface for ISEA3H compaction."""
    parser = argparse.ArgumentParser(description="ISEA3H Compact")
    parser.add_argument(
        "-i",
        "--input",
        type=str,
        required=True,
        help="Input ISEA3H (GeoJSON, Shapefile, CSV, Parquet, or pickled GeoDataFrame .gpd/.geopandas)",
    )
    parser.add_argument("-cellid", "--cellid", type=str, help="ISEA3H ID field")
    parser.add_argument("-f", "--output_format", type=str, default=None, help="Output format (None, csv, geojson, shapefile, gpd, geojson_dict, gpkg, geoparquet)")

    args = parser.parse_args()
    input_data = args.input
    cellid = args.cellid
    output_format = args.output_format
    
    if platform.system() != "Windows":
        print("ISEA3H DGGS conversion is only supported on Windows")
        return
    
    result = isea3hcompact(
        input_data,
        isea3h_id=cellid,
        output_format=output_format,
    )
    
    if output_format is None:
        print(result)
    elif output_format in ["csv", "geojson", "geojson_dict", "shapefile", "gpkg", "geoparquet", "parquet"]:
        if isinstance(input_data, str):
            base = os.path.splitext(os.path.basename(input_data))[0]
            ext_map = {
                "csv": ".csv",
                "geojson": ".geojson",
                "geojson_dict": ".geojson",
                "shapefile": ".shp",
                "gpkg": ".gpkg",
                "parquet": ".parquet",
                "geoparquet": ".parquet",
            }
            ext = ext_map.get(output_format, "")
            output = f"{base}_isea3h_compacted{ext}"
        else:
            output = f"isea3h_compacted{ext_map.get(output_format, '')}"
        print(f"Output written to {output}")
    elif output_format in ["gpd", "geopandas"]:
        print(result)
    else:
        print("ISEA3H compact completed.")

def isea3h_expand(isea3h_ids, resolution):
    """Expands a list of DGGS cells to the target resolution."""
    if platform.system() != "Windows":
        raise RuntimeError("ISEA3H expansion is only supported on Windows.")
    
    expand_cells = []
    for isea3h_id in isea3h_ids:
        isea3h_cell = DggsCell(isea3h_id)
        expand_cells.extend(get_isea3h_cell_children(isea3h_cell, resolution))
    return expand_cells

def isea3hexpand(
    input_data,
    resolution,
    isea3h_id=None,
    output_format=None,
):
    """Expand ISEA3H cells to a target resolution."""
    if platform.system() != "Windows":
        raise RuntimeError("ISEA3H expansion is only supported on Windows.")
    
    if isea3h_id is None:
        isea3h_id = "isea3h"
    
    gdf = process_input_data_compact(input_data, isea3h_id)
    isea3h_ids = gdf[isea3h_id].drop_duplicates().tolist()
    
    if not isea3h_ids:
        print(f"No ISEA3H IDs found in <{isea3h_id}> field.")
        return
    
    try:
        max_res = max(get_isea3h_resolution(isea3h_id) for isea3h_id in isea3h_ids)
        if resolution < max_res:
            print(f"Target expand resolution ({resolution}) must >= {max_res}.")
            return None
        
        isea3h_cells_expand = isea3h_expand(isea3h_ids, resolution)
        isea3h_ids_expand = [cell.get_cell_id() for cell in isea3h_cells_expand]
    except Exception:
        raise Exception("Expand cells failed. Please check your ISEA3H ID field and resolution.")
    
    if not isea3h_ids_expand:
        return None
    
    rows = []
    for isea3h_id_expand in isea3h_ids_expand:
        try:
            cell_polygon = isea3h2geo(isea3h_id_expand)
            cell_resolution = resolution
            num_edges = 6  # ISEA3H cells are hexagonal
            row = geodesic_dggs_to_geoseries(
                "isea3h", isea3h_id_expand, cell_resolution, cell_polygon, num_edges
            )
            rows.append(row)
        except Exception:
            continue
    
    out_gdf = gpd.GeoDataFrame(rows, geometry="geometry",crs="EPSG:4326")
    
    file_formats = ["csv", "geojson", "shapefile", "gpkg", "parquet", "geoparquet"]
    output_name = None
    if output_format in file_formats:
        ext_map = {
            "csv": ".csv",
            "geojson": ".geojson",
            "shapefile": ".shp",
            "gpkg": ".gpkg",
            "parquet": ".parquet",
            "geoparquet": ".parquet",
        }
        ext = ext_map.get(output_format, "")
        if isinstance(input_data, str):
            base = os.path.splitext(os.path.basename(input_data))[0]
            output_name = f"{base}_isea3h_expanded{ext}"
        else:
            output_name = f"isea3h_expanded{ext}"
    
    return convert_to_output_format(out_gdf, output_format, output_name)

def isea3hexpand_cli():
    """Command-line interface for ISEA3H expansion."""
    parser = argparse.ArgumentParser(description="ISEA3H Expand (Uncompact)")
    parser.add_argument(
        "-i",
        "--input",
        type=str,
        required=True,
        help="Input ISEA3H (GeoJSON, Shapefile, CSV, Parquet, or pickled GeoDataFrame .gpd/.geopandas)",
    )
    parser.add_argument(
        "-r",
        "--resolution",
        type=int,
        required=True,
        help="Target ISEA3H resolution to expand to (must be greater than input cells)",
    )
    parser.add_argument("-cellid", "--cellid", type=str, help="ISEA3H ID field")
    parser.add_argument("-f", "--output_format", type=str, default=None, help="Output format (None, csv, geojson, shapefile, gpd, geojson_dict, gpkg, geoparquet)")

    args = parser.parse_args()
    input_data = args.input
    resolution = args.resolution
    cellid = args.cellid
    output_format = args.output_format
    
    if platform.system() != "Windows":
        print("ISEA3H DGGS conversion is only supported on Windows")
        return
    
    result = isea3hexpand(
        input_data,
        resolution,
        isea3h_id=cellid,
        output_format=output_format,
    )
    
    if output_format is None:
        print(result)
    elif output_format in ["csv", "geojson", "geojson_dict", "shapefile", "gpkg", "geoparquet", "parquet"]:
        if isinstance(input_data, str):
            base = os.path.splitext(os.path.basename(input_data))[0]
            ext_map = {
                "csv": ".csv",
                "geojson": ".geojson",
                "geojson_dict": ".geojson",
                "shapefile": ".shp",
                "gpkg": ".gpkg",
                "parquet": ".parquet",
                "geoparquet": ".parquet",
            }
            ext = ext_map.get(output_format, "")
            output = f"{base}_isea3h_expanded{ext}"
        else:
            output = f"isea3h_expanded{ext_map.get(output_format, '')}"
        print(f"Output written to {output}")
    elif output_format in ["gpd", "geopandas"]:
        print(result)
    else:
        print("ISEA3H expand completed.") 