"""
Copyright (C) 2025 Applied Geospatial Research Group.

This script is licensed under the GNU General Public License v3.0.
See <https://gnu.org/licenses/gpl-3.0> for full license details.

Author: Richard Zeng, Maverick Fong

Description:
    This script is part of the BERA Tools.
    Webpage: https://github.com/appliedgrg/beratools

    This file hosts the ground_footprint tool.
"""

import logging
import math
from itertools import chain

import geopandas as gpd
import numpy as np
import pandas as pd
import pyogrio.errors
import shapely.geometry as sh_geom
import shapely.ops as sh_ops

import beratools.core.algo_common as algo_common
import beratools.core.constants as bt_const
import beratools.utility.spatial_common as sp_common
from beratools.core.algo_line_grouping import LineGrouping
from beratools.core.algo_merge_lines import custom_line_merge
from beratools.core.algo_split_with_lines import LineSplitter
from beratools.core.logger import Logger
from beratools.core.tool_base import execute_multiprocessing
from beratools.utility.tool_args import CallMode

log = Logger("ground_footprint", file_level=logging.INFO)
logger = log.get_logger()
print = log.print

FP_FIXED_WIDTH_DEFAULT = 5.0


def prepare_line_args(line_gdf, poly_gdf, n_samples, offset, width_percentile):
    """
    Generate arguments for each line in the GeoDataFrame.

    Args:
        line_gdf
        poly_gdf
        n_samples
        offset

    Returns:
        line_args : list
            row :
            inter_poly :
            n_samples :
            offset :
            width_percentile :

    """
    spatial_index = poly_gdf.sindex
    line_args = []

    for idx in line_gdf.index:
        row = line_gdf.loc[[idx]]
        line = row.geometry.iloc[0]

        # Skip rows where geometry is None
        if line is None:
            print(row)
            continue

        inter_poly = poly_gdf.loc[spatial_index.query(line)]
        if bt_const.BT_GROUP in inter_poly.columns:
            inter_poly = inter_poly[inter_poly[bt_const.BT_GROUP] == row[bt_const.BT_GROUP].values[0]]

        try:
            line_args.append([row, inter_poly, n_samples, offset, width_percentile])
        except Exception as e:
            print(e)

    return line_args


# Calculating Line Widths
def generate_sample_points(line, n_samples=10):
    """
    Generate evenly spaced points along a line.

    Args:
        line (LineString): The line along which to generate points.
        n_samples (int): The number of points to generate (default is 10).

    Returns:
        list:  List of shapely Point objects.

    """
    # TODO: determine line type
    try:
        pts = line.coords
    except Exception as e:  # TODO: check the code
        print(e)
        line = sh_ops.linemerge(line)
        tuple_coord = sh_geom.mapping(line)["coordinates"]
        pts = list(chain(*tuple_coord))

    return [sh_geom.Point(item) for item in pts]


def process_single_line(line_arg):
    row = line_arg[0]
    inter_poly = line_arg[1]
    n_samples = line_arg[2]
    offset = line_arg[3]
    width_percentile = line_arg[4]

    # TODO: deal with case when inter_poly is empty
    try:
        widths, line, perp_lines, perp_lines_original = calculate_average_width(
            row.iloc[0].geometry, inter_poly, offset, n_samples
        )
    except Exception as e:
        print(e)
        return None

    # Calculate the 75th percentile width
    # filter zeros in width array
    arr_filter = [False if math.isclose(i, 0.0) else True for i in widths]
    widths = widths[arr_filter]

    q3_width = FP_FIXED_WIDTH_DEFAULT
    q4_width = FP_FIXED_WIDTH_DEFAULT
    try:
        # TODO: check the code. widths is empty
        if len(widths) > 0:
            q3_width = np.percentile(widths, width_percentile)
            q4_width = np.percentile(widths, 90)
    except Exception as e:
        print(e)

    # Store the 75th percentile width as a new attribute
    row["avg_width"] = q3_width
    row["max_width"] = q4_width

    row["geometry"] = line
    try:
        row["perp_lines"] = perp_lines
        row["perp_lines_original"] = perp_lines_original
    except Exception as e:
        print(e)

    return row


def generate_fixed_width_footprint(line_gdf, max_width=False):
    """
    Create a buffer around each line.

    In the GeoDataFrame using its 'max_width' attribute and
    saves the resulting polygons in a new shapefile.

    Args:
    line_gdf: GeoDataFrame containing LineString with 'max_width' attribute.
    max_width: Use max width or not to produce buffer.

    """
    # Create a new GeoDataFrame with the buffer polygons
    buffer_gdf = line_gdf.copy(deep=True)

    mean_avg_width = line_gdf["avg_width"].mean()
    mean_max_width = line_gdf["max_width"].mean()

    # Use .loc to avoid chained assignment
    line_gdf.loc[line_gdf["avg_width"].isna(), "avg_width"] = mean_avg_width
    line_gdf.loc[line_gdf["max_width"].isna(), "max_width"] = mean_max_width

    line_gdf.loc[line_gdf["avg_width"] == 0.0, "avg_width"] = mean_avg_width
    line_gdf.loc[line_gdf["max_width"] == 0.0, "max_width"] = mean_max_width

    if not max_width:
        print("Using quantile 75% width")
        buffer_gdf["geometry"] = line_gdf.apply(
            lambda row: row.geometry.buffer(row.avg_width / 2) if row.geometry is not None else None,
            axis=1,
        )
    else:
        print("Using quantile 90% + 20% width")
        buffer_gdf["geometry"] = line_gdf.apply(
            lambda row: row.geometry.buffer(row.max_width * 1.2 / 2) if row.geometry is not None else None,
            axis=1,
        )

    return buffer_gdf


def calculate_average_width(line, in_poly, offset, n_samples):
    """Calculate the average width of a polygon perpendicular to the given line."""
    # Smooth the line
    try:
        line = line.simplify(0.1)

        valid_widths = 0
        sample_points = generate_sample_points(line, n_samples=n_samples)
        sample_points_pairs = list(zip(sample_points[:-2], sample_points[1:-1], sample_points[2:]))
        widths = np.zeros(len(sample_points_pairs))
        perp_lines = []
        perp_lines_original = []
    except Exception as e:
        print(e)

    try:
        for i, points in enumerate(sample_points_pairs):
            try:
                perp_line = algo_common.generate_perpendicular_line_precise(points, offset=offset)
                perp_lines_original.append(perp_line)
            except Exception as e:
                print(f"Failed to generate perpendicular at index {i}: {e}")
                perp_lines_original.append(None)
                continue

            try:
                polygon_intersect = in_poly.iloc[in_poly.sindex.query(perp_line)]
                intersections = polygon_intersect.intersection(perp_line)
            except Exception as e:
                print(f"Failed intersection at index {i}: {e}")
                intersections = []

            line_list = []
            for inter in intersections:
                if inter.is_empty:
                    continue

                if isinstance(inter, sh_geom.GeometryCollection):
                    for item in inter.geoms:
                        if isinstance(item, sh_geom.LineString):
                            line_list.append(item)
                elif isinstance(inter, sh_geom.MultiLineString):
                    line_list += list(inter.geoms)
                else:
                    line_list.append(inter)

            perp_lines += line_list

            if isinstance(line_list, sh_geom.GeometryCollection):
                print("Found 2: GeometryCollection")

            for item in line_list:
                widths[i] = max(widths[i], item.length)
                valid_widths += 1

            # Todo: check missing perpendicular lines
            # if len(perp_lines_original) < len(sample_points_pairs):
            #     print(f"Missing perpendicular at index {i}")

    except Exception as e:
        print(f"loop: {e}")

    return (
        widths,
        line,
        sh_geom.MultiLineString(perp_lines),
        sh_geom.MultiLineString(perp_lines_original),
    )


def ground_footprint(
    in_line,
    in_footprint,
    n_samples,
    offset,
    max_width,
    out_footprint,
    in_layer_lc_path="least_cost_path",
    merge_group=True,
    width_percentile=75,
    trim_output=True,
    processes=0,
    call_mode=CallMode.CLI,
    log_level="INFO",
):
    in_file, in_layer = sp_common.decode_file_layer(in_line)
    in_fp_file, in_layer_fp = sp_common.decode_file_layer(in_footprint)
    out_file, out_layer = sp_common.decode_file_layer(out_footprint)

    n_samples = int(n_samples)
    offset = float(offset)
    width_percentile = int(width_percentile)

    # TODO: refactor this code for better line quality check
    line_gdf = gpd.read_file(in_file, layer=in_layer)
    if "fid" in line_gdf.columns:
        line_gdf = line_gdf.rename(columns={"fid": "orig_fid"})
    line_gdf = algo_common.clean_geometries(line_gdf, stage="input")
    line_gdf = line_gdf.reset_index(drop=True)
    if bt_const.BT_GROUP not in line_gdf.columns:
        line_gdf[bt_const.BT_GROUP] = range(1, len(line_gdf) + 1)

    use_least_cost_path = True
    try:
        lc_path_gdf = gpd.read_file(in_file, layer=in_layer_lc_path)
        if "fid" in lc_path_gdf.columns:
            lc_path_gdf = lc_path_gdf.rename(columns={"fid": "orig_fid"})
        lc_path_gdf = algo_common.clean_geometries(lc_path_gdf, stage="input")
        lc_path_gdf = lc_path_gdf.reset_index(drop=True)
    except (ValueError, OSError, pyogrio.errors.DataLayerError):
        print(f"Layer '{in_layer_lc_path}' not found in {in_file}, skipping least cost path logic.")
        use_least_cost_path = False

    if not merge_group:
        line_gdf["geometry"] = line_gdf.geometry.apply(custom_line_merge)
        if use_least_cost_path:
            lc_path_gdf["geometry"] = lc_path_gdf.geometry.apply(custom_line_merge)

    # read footprints and remove holes
    poly_gdf = gpd.read_file(in_fp_file, layer=in_layer_fp)
    if "fid" in poly_gdf.columns:
        poly_gdf = poly_gdf.rename(columns={"fid": "orig_fid"})
    poly_gdf = algo_common.clean_geometries(poly_gdf, stage="input")
    poly_gdf = poly_gdf.reset_index(drop=True)
    poly_gdf["geometry"] = poly_gdf["geometry"].apply(algo_common.remove_holes)

    # merge group and/or split lines at intersections
    merged_line_gdf = line_gdf.copy(deep=True)
    if merge_group:
        lg = LineGrouping(line_gdf, merge_group)
        lg.run_grouping()
        merged_line_gdf = lg.run_line_merge()
    else:
        try:
            lg = LineGrouping(line_gdf, not merge_group)
            lg.run_grouping()
            merged_line_gdf = lg.run_line_merge()
            splitter = LineSplitter(merged_line_gdf)
            splitter.process()
            splitter.save_to_geopackage(
                out_footprint,
                line_layer="split_centerline",
                intersection_layer="inter_points",
                invalid_layer="invalid_splits",
            )

            # least cost path merge and split
            if use_least_cost_path:
                lg_leastcost = LineGrouping(lc_path_gdf, not merge_group)
                lg_leastcost.run_grouping()
                merged_lc_path_gdf = lg_leastcost.run_line_merge()
                splitter_leastcost = LineSplitter(merged_lc_path_gdf)
                splitter_leastcost.process(splitter.intersection_gdf)

                splitter_leastcost.save_to_geopackage(
                    out_footprint,
                    line_layer="split_leastcost",
                )

                lg = LineGrouping(splitter.split_lines_gdf, merge_group)
                lg.run_grouping()
                merged_line_gdf = lg.run_line_merge()
        except ValueError as e:
            print(f"Exception: ground_footprint: {e}")

    # save original merged lines
    merged_line_gdf = algo_common.clean_geometries(
        merged_line_gdf,
        stage="output",
        out_file=out_file,
        layer="rejected_output_ground_merged_lines",
    )
    merged_line_gdf.to_file(out_file, layer="merged_lines_original")

    # prepare line arguments
    line_args = prepare_line_args(merged_line_gdf, poly_gdf, n_samples, offset, width_percentile)

    out_lines = execute_multiprocessing(
        process_single_line, line_args, "Fixed footprint", processes, call_mode
    )
    line_attr = pd.concat(out_lines)

    # Ensure BT_GROUP is present in line_attr
    if bt_const.BT_GROUP not in line_attr.columns:
        raise ValueError("BT_GROUP column is required in line_attr but is missing.")

    # update avg_width and max_width by max value of group
    if not merge_group:
        group_max = (
            line_attr.groupby(bt_const.BT_GROUP).agg({"avg_width": "max", "max_width": "max"}).reset_index()
        )

        # Merge the result back to the original dataframe based on 'group'
        line_attr = line_attr.merge(group_max, on=bt_const.BT_GROUP, suffixes=("", "_max"))

        # Overwrite the original columns directly with the max values
        line_attr["avg_width"] = line_attr["avg_width_max"]
        line_attr["max_width"] = line_attr["max_width_max"]

        # Drop the temporary max columns
        line_attr.drop(columns=["avg_width_max", "max_width_max"], inplace=True)

    # create fixed width footprint (always assign buffer_gdf)
    print("Step: Generating fixed width footprints")
    buffer_gdf = generate_fixed_width_footprint(line_attr, max_width=max_width)

    # reserve all layers for output
    perp_lines_gdf = buffer_gdf.copy(deep=True)
    perp_lines_original_gdf = buffer_gdf.copy(deep=True)

    # Save untrimmed fixed width footprint
    buffer_gdf = buffer_gdf.drop(columns=["perp_lines"])
    buffer_gdf = buffer_gdf.drop(columns=["perp_lines_original"])
    buffer_gdf = buffer_gdf.set_crs(perp_lines_gdf.crs, allow_override=True)
    buffer_gdf.reset_index(inplace=True, drop=True)

    untrimmed_footprint = "untrimmed_footprint"
    buffer_gdf = algo_common.clean_geometries(
        buffer_gdf,
        stage="output",
        out_file=out_file,
        layer="rejected_output_ground_untrimmed_footprint",
    )
    buffer_gdf.to_file(out_file, layer=untrimmed_footprint)
    print(f"Untrimmed fixed width footprint saved as '{untrimmed_footprint}'")

    # trim lines and footprints
    if trim_output:
        lg.run_cleanup(buffer_gdf)

        # Ensure only polygons are saved in clean_footprint
        def ensure_polygons(gdf, buffer_width=0.01):
            gdf["geometry"] = gdf["geometry"].apply(
                lambda geom: geom.buffer(buffer_width)
                if geom.geom_type in ["LineString", "MultiLineString"]
                else geom
            )
            gdf = gdf[gdf.geometry.type.isin(["Polygon", "MultiPolygon"])]
            return gdf

        # Patch: after trimming, ensure polygons in clean_footprint layer
        if hasattr(lg, "merged_lines_trimmed") and lg.merged_lines_trimmed is not None:
            lg.merged_lines_trimmed = ensure_polygons(lg.merged_lines_trimmed)
            print("Step: Saving trimmed outputs")
            lg.save_file(out_file, out_layer)
        else:
            print("Skipping line and footprint trimming per user option.")

    # perpendicular lines
    layer = "perp_lines"
    out_aux_gpkg = algo_common.get_aux_path(out_file)
    perp_lines_gdf = perp_lines_gdf.set_geometry("perp_lines")
    perp_lines_gdf = perp_lines_gdf.drop(columns=["perp_lines_original"])
    perp_lines_gdf = perp_lines_gdf.drop(columns=["geometry"])
    perp_lines_gdf = perp_lines_gdf.set_crs(buffer_gdf.crs, allow_override=True)
    perp_lines_gdf.to_file(out_aux_gpkg, layer=layer)

    layer = "perp_lines_original"
    perp_lines_original_gdf = perp_lines_original_gdf.set_geometry("perp_lines_original")
    perp_lines_original_gdf = perp_lines_original_gdf.drop(columns=["perp_lines"])
    perp_lines_original_gdf = perp_lines_original_gdf.drop(columns=["geometry"])
    perp_lines_original_gdf = perp_lines_original_gdf.set_crs(buffer_gdf.crs, allow_override=True)
    perp_lines_original_gdf.to_file(out_aux_gpkg, layer=layer)

    layer = "centerline_simplified"
    # Drop perp_lines_original column if present to avoid export warnings
    if "perp_lines_original" in line_attr.columns:
        line_attr = line_attr.drop(columns=["perp_lines_original"])
    line_attr = line_attr.drop(columns="perp_lines")
    line_attr.to_file(out_aux_gpkg, layer=layer)

    # save footprints without holes
    poly_gdf.to_file(out_aux_gpkg, layer="footprint_no_holes")
    print("Step: Finished fixed width footprint tool")


if __name__ == "__main__":
    import time

    from beratools.utility.tool_args import compose_tool_kwargs

    start_time = time.time()
    kwargs = compose_tool_kwargs("ground_footprint")
    ground_footprint(**kwargs)
    print("Elapsed time: {}".format(time.time() - start_time))
