"""
Helper functions for working with OPR radar datasets.

This module provides utility functions for common operations on radar datasets,
including merging frames from the same flight segment, generating citations, and
finding crossover points between flight lines.
"""

import warnings
from typing import Iterable, Union

import geopandas as gpd
import xarray as xr

from xopr.util import get_ror_display_name, merge_dicts_no_conflicts


def merge_frames(frames: Iterable[xr.Dataset]) -> Union[list[xr.Dataset], xr.Dataset]:
    """
    Merge radar frames by segment.

    Frames from the same segment are concatenated along slow_time and sorted.

    Parameters
    ----------
    frames : Iterable[xr.Dataset]
        Radar frame Datasets to merge.

    Returns
    -------
    xr.Dataset or list[xr.Dataset]
        Single merged Dataset if all frames are from one segment, otherwise
        list of Datasets (one per segment).
    """
    segments = {}

    mergable_keys = ['source_url', 'collection', 'data_product', 'granule']

    for frame in frames:
        # Get segment path from frame attributes
        granule = frame.attrs.get('granule')
        date, segment_id, frame_id = granule.split('_')
        segment_path = f"{date}_{segment_id}"

        if not segment_path:
            warnings.warn("Frame missing 'granule' attribute or it was not in the expected format, skipping.", UserWarning)
            continue

        if segment_path not in segments:
            segments[segment_path] = []

        segments[segment_path].append(frame)

    # Merge frames for each segment
    merged_segments = []
    for segment_path, segment_frames in segments.items():
        merged_segment = xr.concat(segment_frames, dim='slow_time', combine_attrs=merge_dicts_no_conflicts).sortby('slow_time')

        for k in mergable_keys:
            if k not in merged_segment.attrs:
                merged_segment.attrs[k] = set([v for v in {f.attrs.get(k) for f in segment_frames} if v is not None])

        merged_segments.append(merged_segment)

    if len(merged_segments) == 1:
        return merged_segments[0]
    else:
        return merged_segments

def generate_citation(ds : xr.Dataset) -> str:
        """
        Generate citation and acknowledgment text for dataset.

        Parameters
        ----------
        ds : xr.Dataset
            Dataset with doi, ror, and funder_text attributes.

        Returns
        -------
        str
            Formatted citation string including data citation, DOI, funding
            acknowledgments, and OPR toolbox citation.
        """

        citation_string = ""
        any_citation_info = False

        citation_string += "== Data Citation ==\n"

        if 'ror' in ds.attrs and ds.attrs['ror']:
            any_citation_info = True
            if isinstance(ds.attrs['ror'], (set, list)):
                institution_name = ', '.join([get_ror_display_name(ror) for ror in ds.attrs['ror']])
            else:
                institution_name = get_ror_display_name(ds.attrs['ror'])

            citation_string += f"This data was collected by {institution_name}.\n"

        if 'doi' in ds.attrs and ds.attrs['doi']:
            any_citation_info = True
            citation_string += f"Please cite the dataset DOI: https://doi.org/{ds.attrs['doi']}\n"

        if 'funder_text' in ds.attrs and ds.attrs['funder_text']:
            any_citation_info = True
            citation_string += f"Please include the following funder acknowledgment:\n{ds.attrs['funder_text']}\n"

        if not any_citation_info:
            citation_string += "No specific citation information was retrieved for this dataset. By default, please cite:\n"
            citation_string += "CReSIS. 2024. REPLACE_WITH_RADAR_NAME Data, Lawrence, Kansas, USA. Digital Media. http://data.cresis.ku.edu/."

        # Add general OPR Toolbox citation
        citation_string += "\n== Processing Citation ==\n"
        citation_string += "Data was processed using the Open Polar Radar (OPR) Toolbox: https://doi.org/10.5281/zenodo.5683959\n"
        citation_string += "Please cite the OPR Toolbox as:\n"
        citation_string += "Open Polar Radar. (2024). opr (Version 3.0.1) [Computer software]. https://gitlab.com/openpolarradar/opr/. https://doi.org/10.5281/zenodo.5683959\n"
        citation_string += "And include the following acknowledgment:\n"
        citation_string += "We acknowledge the use of software from Open Polar Radar generated with support from the University of Kansas, NASA grants 80NSSC20K1242 and 80NSSC21K0753, and NSF grants OPP-2027615, OPP-2019719, OPP-1739003, IIS-1838230, RISE-2126503, RISE-2127606, and RISE-2126468.\n"

        return citation_string if citation_string else "No citation information available for this dataset."


def find_intersections(gdf : gpd.GeoDataFrame, remove_self_intersections: bool = True,
                        remove_adjacent_intersections: bool = True,
                        calculate_crossing_angles: bool = False) -> gpd.GeoDataFrame:
    """
    Find crossover points between radar flight lines.

    Parameters
    ----------
    gdf : gpd.GeoDataFrame
        GeoDataFrame of STAC items with flight line geometries.
    remove_self_intersections : bool, optional
        If True, exclude geometry intersecting with itself.
    remove_adjacent_intersections : bool, optional
        If True, exclude consecutive frames in same flight line.
    calculate_crossing_angles : bool, optional
        If True, calculate and include crossing angles in degrees.

    Returns
    -------
    gpd.GeoDataFrame
        GeoDataFrame with columns for intersecting pairs (id_1, id_2),
        intersection_geometry, and optionally crossing_angle.
    """

    tmp_df = gdf.reset_index()
    tmp_df['geom'] = tmp_df.geometry
    intersections = gpd.sjoin(tmp_df, tmp_df, how='inner', predicate='intersects', lsuffix='1', rsuffix='2')

    if remove_self_intersections:
        intersections = intersections[intersections['id_1'] != intersections['id_2']]

    intersections['intersection_geometry'] = intersections.apply(lambda row: row['geom_1'].intersection(row['geom_2']), axis=1)
    intersections.set_geometry('intersection_geometry', inplace=True, crs=gdf.crs)
    intersections = intersections.drop_duplicates(subset=['intersection_geometry'])
    intersections = intersections.explode(index_parts=True).reset_index(drop=True)

    intersections_tmp = intersections[['id_1', 'id_2', 'intersection_geometry', 'collection_1', 'collection_2', 'geom_1', 'geom_2']].copy()

    for k in ['opr:date', 'opr:segment', 'opr:frame']:
        intersections_tmp[f'{k}_1'] = intersections['properties_1'].apply(lambda x: x[k])
        intersections_tmp[f'{k}_2'] = intersections['properties_2'].apply(lambda x: x[k])

    intersections = intersections_tmp

    if remove_adjacent_intersections:
        intersections = intersections[
            (intersections['opr:date_1'] != intersections['opr:date_2']) |
            (intersections['opr:segment_1'] != intersections['opr:segment_2']) |
            ((intersections['opr:frame_1'] != (intersections['opr:frame_2'] + 1)) &
            (intersections['opr:frame_1'] != (intersections['opr:frame_2'] - 1)))
        ]

    if calculate_crossing_angles:
        intersections['crossing_angle'] = intersections.apply(
            lambda row: _calculate_crossing_angle(row['geom_1'], row['geom_2'], row['intersection_geometry']),
            axis=1
        )

    return intersections

def _calculate_crossing_angle(line1, line2, intersection_point):
    """
    Calculate the crossing angle between two lines at their intersection point.

    Parameters
    ----------
    line1 : shapely.geometry.LineString
        The first line.
    line2 : shapely.geometry.LineString
        The second line.
    intersection_point : shapely.geometry.Point
        The point of intersection between the two lines.

    Returns
    -------
    float
        The crossing angle in degrees.
    """
    import numpy as np

    def get_line_angle(line, point):
        # Get the nearest point on the line to the intersection point
        nearest_point = line.interpolate(line.project(point))
        # Get a small segment of the line around the nearest point to calculate the angle
        buffer_distance = 1e-6  # Small distance to create a segment
        start_point = line.interpolate(max(0, line.project(nearest_point) - buffer_distance))
        end_point = line.interpolate(min(line.length, line.project(nearest_point) + buffer_distance))

        # Calculate the angle of the segment
        delta_x = end_point.x - start_point.x
        delta_y = end_point.y - start_point.y
        angle = np.arctan2(delta_y, delta_x)
        return angle

    angle1 = get_line_angle(line1, intersection_point)
    angle2 = get_line_angle(line2, intersection_point)

    # Calculate the absolute difference in angles and convert to degrees
    angle_diff = np.abs(angle1 - angle2)
    crossing_angle = np.degrees(angle_diff)

    # Ensure the angle is between 0 and 180 degrees
    if crossing_angle > 180:
        crossing_angle = 360 - crossing_angle

    if crossing_angle > 90:
        crossing_angle = 180 - crossing_angle

    return crossing_angle
