import unittest
import numpy as np
from osgeo import gdal, ogr, osr
from geodata_pipeline.river_erosion import _build_order_lines, _tile_geotransform
from geodata_pipeline.config import RiverErosionProfileConfig

class TestRiverContinuity(unittest.TestCase):
    def setUp(self):
        # Create a simple river line crossing x=1000
        # River flows from (500, 500) to (1500, 500)
        self.srs = osr.SpatialReference()
        self.srs.ImportFromEPSG(25832)
        
        self.river_geom = ogr.Geometry(ogr.wkbLineString)
        self.river_geom.AddPoint(500.0, 500.0)
        self.river_geom.AddPoint(1500.0, 500.0)
        
        # Mock Layer
        self.driver = ogr.GetDriverByName("Memory")
        self.ds = self.driver.CreateDataSource("wrk")
        self.layer = self.ds.CreateLayer("rivers", srs=self.srs)
        
        field_defn = ogr.FieldDefn("ORD_STRA", ogr.OFTInteger)
        self.layer.CreateField(field_defn)
        
        feat = ogr.Feature(self.layer.GetLayerDefn())
        feat.SetGeometry(self.river_geom)
        feat.SetField("ORD_STRA", 5)
        self.layer.CreateFeature(feat)
        
        self.profile = RiverErosionProfileConfig() # Defaults
        self.transform = osr.CoordinateTransformation(self.srs, self.srs) # Identity

    def test_continuity_at_boundary(self):
        # Tile A: 0-1000. Mosaic (3x3): -1000 to 2000.
        bounds_a = (-1000.0, -1000.0, 2000.0, 2000.0)
        # Resolution 1m -> 3000 x 3000 pixels
        w_a = 3001 # 3 tiles * 1000m + 1
        h_a = 3001
        
        # Tile B: 1000-2000. Mosaic (3x3): 0 to 3000.
        bounds_b = (0.0, -1000.0, 3000.0, 2000.0)
        w_b = 3001
        h_b = 3001
        
        # Rasterize A
        res_a = _build_order_lines(
            [(self.layer, "rivers")],
            bounds_a, w_a, h_a,
            self.profile, self.transform, self.transform, self.srs
        )
        
        # Rasterize B
        res_b = _build_order_lines(
            [(self.layer, "rivers")],
            bounds_b, w_b, h_b,
            self.profile, self.transform, self.transform, self.srs
        )
        
        # Compare at X=1000
        # For A: X=1000 is at index 2000 (0 -> 1000 is 1000px, -1000 start)
        # origin_a = -1000. 1000 - (-1000) = 2000.
        col_a = 2000
        
        # For B: X=1000 is at index 1000
        # origin_b = 0. 1000 - 0 = 1000.
        col_b = 1000
        
        # Compare a vertical slice at the seam
        slice_a = res_a[:, col_a]
        slice_b = res_b[:, col_b]
        
        # They should be identical
        diff = np.abs(slice_a - slice_b)
        max_diff = np.max(diff)
        
        self.assertEqual(max_diff, 0.0, f"Seam mismatch found! Max diff: {max_diff}")

if __name__ == "__main__":
    unittest.main()