import math
import os
import numpy as np
from datetime import datetime
from collections import defaultdict

import shapely
from shapely.geometry import LineString
from shapely.ops import unary_union

from qgis.PyQt.QtCore import QVariant
from qgis.core import (
    QgsFeature, QgsField, QgsFields, QgsGeometry, QgsPointXY,
    QgsProcessing, QgsProcessingAlgorithm, QgsProcessingParameterNumber,
    QgsProcessingParameterString, QgsProcessingParameterVectorLayer,
    QgsProcessingParameterRasterDestination, QgsProcessingParameterRasterLayer,
    QgsProcessingParameterFeatureSink, QgsProcessingParameterBoolean, QgsWkbTypes, 
    QgsFeatureRequest, QgsFeatureSink
)
from osgeo import gdal

from .common import DEFAULT_REF_DATE, parse_ref_date

class HovmollerUltraFinalMatch(QgsProcessingAlgorithm):
    INPUT_LINES = 'INPUT_LINES'
    INPUT_BASELINE = 'INPUT_BASELINE'
    DATE_FIELD = 'DATE_FIELD'
    LONGI_FIELD = 'LONGI_FIELD'
    SEARCH_DIST = 'SEARCH_DIST'
    SMOOTH_WINDOW = 'SMOOTH_WINDOW'
    PIXEL_SIZE = 'PIXEL_SIZE'
    INVERT_Y = 'INVERT_Y'
    LONGI_SCALE = 'LONGI_SCALE'
    REF_RASTER = 'REF_RASTER'
    REF_DATE = 'REF_DATE'
    GRID_YEAR_INTERVAL = 'GRID_YEAR_INTERVAL'
    GRID_LONGI_INTERVAL = 'GRID_LONGI_INTERVAL'

    OUTPUT_RAW_PTS = 'OUTPUT_RAW_PTS'
    OUTPUT_GRID_PTS = 'OUTPUT_GRID_PTS'
    OUTPUT_WIDTH_RASTER = 'OUTPUT_WIDTH_RASTER'
    OUTPUT_CHANGE_RASTER = 'OUTPUT_CHANGE_RASTER'
    OUTPUT_AXIS_GRID = 'OUTPUT_AXIS_GRID'
    OUTPUT_AXIS_LABELS = 'OUTPUT_AXIS_LABELS'

    def name(self): return 'hovmoller_step2_v23'
    def displayName(self): return '1. Analysis: Hovmöller Diagram (Width & Change)'
    def group(self): return 'Shoreline Analysis'
    def groupId(self): return 'shoreline'
    def createInstance(self): return HovmollerUltraFinalMatch()

    def shortHelpString(self):
        return (
            "<h3>Step 2 — Hovmöller diagram</h3>"
            "<p>Casts transects from the baseline against each dated shoreline and "
            "builds the width and change rasters plus the raw and grid point layers. "
            "The reference date is the epoch for the integer day index.</p>"
            "<p>Also outputs an axis grid overlay: dashed lines at every Nth year "
            "(time axis) and every Nth longi segment (along-shore axis), with a "
            "label point layer (years along the bottom, segment numbers along the "
            "left edge) so the diagram can be read directly.</p>"
        )

    def initAlgorithm(self, config=None):
        self.addParameter(QgsProcessingParameterVectorLayer(self.INPUT_LINES, 'Clipped Shorelines', [QgsProcessing.TypeVectorLine]))
        self.addParameter(QgsProcessingParameterVectorLayer(self.INPUT_BASELINE, 'Clipped Baseline', [QgsProcessing.TypeVectorLine]))
        self.addParameter(QgsProcessingParameterString(self.DATE_FIELD, 'Clean Date Field (date_clean)', defaultValue='date_clean'))
        self.addParameter(QgsProcessingParameterString(self.LONGI_FIELD, 'Longi Field', defaultValue='longi'))
        self.addParameter(QgsProcessingParameterNumber(self.SEARCH_DIST, 'Transect Length (m)', defaultValue=500.0))
        self.addParameter(QgsProcessingParameterNumber(self.SMOOTH_WINDOW, 'Transect Direction Smoothing', defaultValue=10))
        self.addParameter(QgsProcessingParameterNumber(self.PIXEL_SIZE, 'Pixel Size', defaultValue=10.0))
        self.addParameter(QgsProcessingParameterBoolean(self.INVERT_Y, 'Invert Y axis for Hovmöller diagram', defaultValue=False))
        self.addParameter(QgsProcessingParameterNumber(self.LONGI_SCALE, 'Longi Scale Factor (Stretch/Compress)', type=QgsProcessingParameterNumber.Double, defaultValue=1.0))
        self.addParameter(QgsProcessingParameterRasterLayer(self.REF_RASTER, 'Reference Raster (Optional)', optional=True))
        self.addParameter(QgsProcessingParameterString(self.REF_DATE, 'Reference Date (epoch for day index, YYYY-MM-DD)', defaultValue=DEFAULT_REF_DATE))
        self.addParameter(QgsProcessingParameterNumber(self.GRID_YEAR_INTERVAL, 'Axis Grid: Year Interval', defaultValue=1, minValue=1))
        self.addParameter(QgsProcessingParameterNumber(self.GRID_LONGI_INTERVAL, 'Axis Grid: Longi Interval (segments)', defaultValue=10, minValue=1))

        self.addParameter(QgsProcessingParameterFeatureSink(self.OUTPUT_RAW_PTS, 'Result: Intersection Points (Real)'))
        self.addParameter(QgsProcessingParameterFeatureSink(self.OUTPUT_GRID_PTS, 'Result: Graph Points (Grid)'))
        self.addParameter(QgsProcessingParameterRasterDestination(self.OUTPUT_WIDTH_RASTER, 'Hovmöller Raster width'))
        self.addParameter(QgsProcessingParameterRasterDestination(self.OUTPUT_CHANGE_RASTER, 'Hovmöller Raster change'))
        self.addParameter(QgsProcessingParameterFeatureSink(self.OUTPUT_AXIS_GRID, 'Hovmöller Axis Grid (Lines)', optional=True, createByDefault=True))
        self.addParameter(QgsProcessingParameterFeatureSink(self.OUTPUT_AXIS_LABELS, 'Hovmöller Axis Labels (Points)', optional=True, createByDefault=True))

    def processAlgorithm(self, parameters, context, feedback):
        shoreline_lyr = self.parameterAsVectorLayer(parameters, self.INPUT_LINES, context)
        baseline_lyr = self.parameterAsVectorLayer(parameters, self.INPUT_BASELINE, context)
        date_f = self.parameterAsString(parameters, self.DATE_FIELD, context)
        longi_f = self.parameterAsString(parameters, self.LONGI_FIELD, context)
        dist_max = self.parameterAsDouble(parameters, self.SEARCH_DIST, context)
        smooth_n = self.parameterAsInt(parameters, self.SMOOTH_WINDOW, context)
        px_size = self.parameterAsDouble(parameters, self.PIXEL_SIZE, context)
        invert_y = self.parameterAsBool(parameters, self.INVERT_Y, context)
        longi_scale = self.parameterAsDouble(parameters, self.LONGI_SCALE, context)
        ref_layer = self.parameterAsRasterLayer(parameters, self.REF_RASTER, context)
        grid_year_int = self.parameterAsInt(parameters, self.GRID_YEAR_INTERVAL, context)
        grid_longi_int = self.parameterAsInt(parameters, self.GRID_LONGI_INTERVAL, context)

        output_width_path = self.parameterAsOutputLayer(parameters, self.OUTPUT_WIDTH_RASTER, context)
        output_change_path = self.parameterAsOutputLayer(parameters, self.OUTPUT_CHANGE_RASTER, context)

        ref_date = parse_ref_date(self.parameterAsString(parameters, self.REF_DATE, context))

        # 1. Collect shorelines
        sh_geoms = {}
        temp_dict = defaultdict(list)
        skipped = 0
        for f in shoreline_lyr.getFeatures():
            try:
                dt = datetime.strptime(str(f[date_f]).strip(), '%Y%m%d')
                geom = shapely.from_wkb(bytes(f.geometry().asWkb()))
            except Exception:  # feature has an unparseable date or invalid geometry
                skipped += 1
                continue
            temp_dict[dt].append(geom)
        if skipped:
            feedback.pushWarning(
                f"Skipped {skipped} shoreline feature(s) with invalid '{date_f}' date or geometry.")
        for dt, gs in temp_dict.items():
            sh_geoms[dt] = unary_union(gs)
        
        if not sh_geoms: raise Exception("No valid dates found.")
        sorted_dates = sorted(sh_geoms.keys())
        day_indices = [(dt - ref_date).days for dt in sorted_dates]
        min_day, max_day = min(day_indices), max(day_indices)

        # 2. Prepare baseline
        req = QgsFeatureRequest().addOrderBy(longi_f, True)
        base_features = list(baseline_lyr.getFeatures(req))
        base_data = []
        for f in base_features:
            v = list(f.geometry().vertices())
            if len(v) < 2: continue
            p1, p2 = v[0], v[-1]
            dx, dy = p2.x() - p1.x(), p2.y() - p1.y()
            L = math.hypot(dx, dy)
            if L == 0: continue
            base_data.append({'cx':(p1.x()+p2.x())/2, 'cy':(p1.y()+p2.y())/2, 'nx':-dy/L, 'ny':dx/L, 'longi':int(f[longi_f])})

        longi_vals = [d['longi'] for d in base_data]
        min_longi, max_longi = min(longi_vals), max(longi_vals)

        # 3. Matrix dimensions
        width = int(math.ceil((max_day - min_day) / px_size)) + 1
        # Масштабуємо висоту матриці згідно з longi_scale
        height = int(math.ceil(((max_longi - min_longi) * longi_scale) / px_size)) + 1
        data_matrix = np.full((height, width), np.nan, dtype=np.float32)

        # 4. Origin
        if ref_layer:
            ext = ref_layer.extent()
            X0, Y0, crs_wkt = float(ext.xMinimum()), float(ext.yMaximum()), ref_layer.crs().toWkt()
        else:
            X0, Y0, crs_wkt = 0.0, 0.0, baseline_lyr.crs().toWkt()
        gt = [X0, px_size, 0.0, Y0, 0.0, -px_size]

        # 5. Sinks
        fields = QgsFields()
        fields.append(QgsField('longi', QVariant.Int)); fields.append(QgsField('days', QVariant.Int)); fields.append(QgsField('dist_m', QVariant.Double))
        (raw_sink, raw_id) = self.parameterAsSink(parameters, self.OUTPUT_RAW_PTS, context, fields, QgsWkbTypes.Point, baseline_lyr.crs())
        (grid_sink, grid_id) = self.parameterAsSink(parameters, self.OUTPUT_GRID_PTS, context, fields, QgsWkbTypes.Point, baseline_lyr.crs())

        # 6. Core loop
        for i, d in enumerate(base_data):
            if feedback.isCanceled(): break
            win = base_data[max(0, i-smooth_n) : min(len(base_data), i+smooth_n+1)]
            snx, sny = sum(w['nx'] for w in win), sum(w['ny'] for w in win)
            L_avg = math.hypot(snx, sny)
            pdx, pdy = (snx/L_avg, sny/L_avg) if L_avg > 0 else (d['nx'], d['ny'])

            for dt, day in zip(sorted_dates, day_indices):
                res = self.fast_cast(d['cx'], d['cy'], pdx, pdy, sh_geoms[dt], dist_max)
                if res:
                    hit_pt, dist_val = res
                    f_raw = QgsFeature(fields)
                    f_raw.setGeometry(QgsGeometry.fromPointXY(QgsPointXY(hit_pt.x, hit_pt.y)))
                    f_raw.setAttributes([d['longi'], day, dist_val])
                    raw_sink.addFeature(f_raw, QgsFeatureSink.FastInsert)
                    
                    col = int(round((day - min_day) / px_size))
                    
                    # Розрахунок рядка з урахуванням longi_scale
                    diff = (d['longi'] - min_longi) if invert_y else (max_longi - d['longi'])
                    row = int(round((diff * longi_scale) / px_size))
                    
                    if 0 <= col < width and 0 <= row < height:
                        data_matrix[row, col] = dist_val
                        pt_x = X0 + col * px_size + (px_size / 2.0)
                        pt_y = Y0 - row * px_size - (px_size / 2.0)
                        f_grid = QgsFeature(fields)
                        f_grid.setGeometry(QgsGeometry.fromPointXY(QgsPointXY(pt_x, pt_y)))
                        f_grid.setAttributes([d['longi'], day, dist_val])
                        grid_sink.addFeature(f_grid, QgsFeatureSink.FastInsert)
            
            if i % 100 == 0: feedback.setProgress(int(i / len(base_data) * 100))

        # 7. Interpolation of Width Matrix
        # 7.1 Горизонтальна інтерполяція (по рядках / у часі)
        for r in range(height):
            row_data = data_matrix[r, :]
            mask = np.isnan(row_data)
            if not np.all(mask):
                valid_idx = np.where(~mask)[0]
                if len(valid_idx) > 1:
                    row_data[mask] = np.interp(np.where(mask)[0], valid_idx, row_data[valid_idx], left=np.nan, right=np.nan)
                    data_matrix[r, :] = row_data

        # 7.2 Вертикальна інтерполяція (по стовпцях / по відстані), щоб заповнити "дірки" від масштабування
        for c in range(width):
            col_data = data_matrix[:, c]
            mask = np.isnan(col_data)
            if not np.all(mask):
                valid_idx = np.where(~mask)[0]
                if len(valid_idx) > 1:
                    col_data[mask] = np.interp(np.where(mask)[0], valid_idx, col_data[valid_idx], left=np.nan, right=np.nan)
                    data_matrix[:, c] = col_data

        # 8. Calculate Change Matrix based on interpolated width
        change_matrix = np.full((height, width), np.nan, dtype=np.float32)
        for r in range(height):
            row_data = data_matrix[r, :]
            valid_cells = np.where(~np.isnan(row_data))[0]
            if len(valid_cells) > 0:
                initial_value = row_data[valid_cells[0]]
                change_matrix[r, :] = row_data - initial_value

        # 9. GDAL Export
        self.save_raster(output_width_path, data_matrix, width, height, gt, crs_wkt)
        self.save_raster(output_change_path, change_matrix, width, height, gt, crs_wkt)

        # 10. Axis grid overlay: years along X (time), longi segments along Y.
        # Diagram coordinate mapping: x = X0 + (day - min_day); the pixel size
        # cancels out (col * px_size), so 1 CRS unit == 1 day. Likewise
        # y = Y0 - diff * longi_scale for a longi difference of `diff` segments.
        axis_fields = QgsFields()
        axis_fields.append(QgsField('axis', QVariant.String))
        axis_fields.append(QgsField('label', QVariant.String))
        axis_fields.append(QgsField('value', QVariant.Int))

        axis_sink = axis_id = lab_sink = lab_id = None
        if parameters.get(self.OUTPUT_AXIS_GRID) is not None:
            (axis_sink, axis_id) = self.parameterAsSink(parameters, self.OUTPUT_AXIS_GRID, context,
                                                        axis_fields, QgsWkbTypes.LineString, baseline_lyr.crs())
        if parameters.get(self.OUTPUT_AXIS_LABELS) is not None:
            (lab_sink, lab_id) = self.parameterAsSink(parameters, self.OUTPUT_AXIS_LABELS, context,
                                                      axis_fields, QgsWkbTypes.Point, baseline_lyr.crs())

        if axis_sink is not None or lab_sink is not None:
            x_min, x_max = X0, X0 + width * px_size
            y_min, y_max = Y0 - height * px_size, Y0

            def _emit(sink, geom, axis, value):
                if sink is None:
                    return
                f = QgsFeature(axis_fields)
                f.setGeometry(geom)
                f.setAttributes([axis, str(value), int(value)])
                sink.addFeature(f, QgsFeatureSink.FastInsert)

            # Vertical grid lines at January 1st of every Nth year; label at the bottom edge.
            for year in range(sorted_dates[0].year, sorted_dates[-1].year + 1, grid_year_int):
                d = (datetime(year, 1, 1) - ref_date).days
                if d < min_day or d > max_day:
                    continue
                x = X0 + (d - min_day)
                _emit(axis_sink, QgsGeometry.fromPolylineXY(
                    [QgsPointXY(x, y_min), QgsPointXY(x, y_max)]), 'time', year)
                _emit(lab_sink, QgsGeometry.fromPointXY(QgsPointXY(x, y_min)), 'time', year)

            # Horizontal grid lines at every Nth longi segment; label at the left edge.
            start_l = ((min_longi + grid_longi_int - 1) // grid_longi_int) * grid_longi_int
            for L in range(int(start_l), int(max_longi) + 1, grid_longi_int):
                diff = (L - min_longi) if invert_y else (max_longi - L)
                y = Y0 - diff * longi_scale
                _emit(axis_sink, QgsGeometry.fromPolylineXY(
                    [QgsPointXY(x_min, y), QgsPointXY(x_max, y)]), 'longi', L)
                _emit(lab_sink, QgsGeometry.fromPointXY(QgsPointXY(x_min, y)), 'longi', L)

        results = {
            self.OUTPUT_RAW_PTS: raw_id,
            self.OUTPUT_GRID_PTS: grid_id,
            self.OUTPUT_WIDTH_RASTER: str(output_width_path),
            self.OUTPUT_CHANGE_RASTER: str(output_change_path)
        }
        if axis_id:
            results[self.OUTPUT_AXIS_GRID] = axis_id
        if lab_id:
            results[self.OUTPUT_AXIS_LABELS] = lab_id
        return results

    def save_raster(self, path, data, w, h, gt, crs):
        driver = gdal.GetDriverByName('GTiff')
        ds = driver.Create(str(path), w, h, 1, gdal.GDT_Float32)
        ds.SetGeoTransform(gt); ds.SetProjection(crs)
        band = ds.GetRasterBand(1)
        band.WriteArray(data); band.SetNoDataValue(-9999)
        ds = None 

    def fast_cast(self, cx, cy, pdx, pdy, shoreline, max_d):
        line = LineString([(cx - max_d*pdx, cy - max_d*pdy), (cx + max_d*pdx, cy + max_d*pdy)])
        inter = line.intersection(shoreline)
        if inter.is_empty: return None
        pts = []
        if inter.geom_type == 'Point': pts = [inter]
        elif hasattr(inter, 'geoms'): pts = [g for g in inter.geoms if g.geom_type == 'Point']
        if not pts: return None
        valid = []
        for p in pts:
            dot = (p.x - cx) * pdx + (p.y - cy) * pdy
            if dot >= 0: valid.append((p, dot))
        if not valid: return None
        best_p, _ = min(valid, key=lambda x: (x[0].x - cx)**2 + (x[0].y - cy)**2)
        dist_sq = (best_p.x - cx)**2 + (best_p.y - cy)**2
        if dist_sq > (max_d * 0.8)**2: return None
        return best_p, math.sqrt(dist_sq)