import pandas as pd
# loading in village boundaries
from geopandas import *
lao_adm_boundaries = gpd.read_file('lao_bnd_village_cde2015.shp',  encoding='cp874').to_crs(epsg=4326)

df = pd.read_csv('lao_pop_projections.csv')
# restrict to adm level 2
df = df[df['ADM_LEVEL'] == 2]
# extract columns GEO_MATCH, ADM2_NAME, BTOTL_2020, BTOTL_2024
df = df[['GEO_MATCH', 'ADM2_NAME', 'BTOTL_2020', 'BTOTL_2024']]

# create a new column DCODE in df, by splitting by _, taking the second and third, removing leading zeros, and concatenating without separator
df['DCODE'] = df['GEO_MATCH'].str.split('_').str[1].str.lstrip('0') + df['GEO_MATCH'].str.split('_').str[2]
# convert to integer
df['DCODE'] = df['DCODE'].astype(int)

# merge df with lao_adm_boundaries restricted to columns DCODE, DNAME on DCODE
lao_adm_boundaries_merge = lao_adm_boundaries[['DCODE', 'DNAME']]
# take unique rows
# number of rows before
lao_adm_boundaries_merge = lao_adm_boundaries_merge.drop_duplicates()
print('Number of rows before:', len(lao_adm_boundaries_merge))
df_merged = df.merge(lao_adm_boundaries_merge, on='DCODE', how='left')

# one value has a mistake
# change the value of DCODE in df where DCODE is 1011 to 1013 and redo the merge
df.loc[df['DCODE'] == 1011, 'DCODE'] = 1013
df_merged = df.merge(lao_adm_boundaries_merge, on='DCODE')

# list all files in the popStructure directory
import os
import rasterio
import numpy as np
from rasterio.mask import mask

filenames = os.listdir('popStructure')

for filename in filenames:
    if filename.endswith('.tif'):
        column_name = os.path.splitext(filename)[0]
        print("Processing", column_name)
        dataset = rasterio.open('popStructure/' + filename)
        print("pop raster CRS:", dataset.crs)
        
        # Ensure lao_adm_boundaries is in the same CRS as the raster
        lao_adm_boundaries = lao_adm_boundaries.to_crs(dataset.crs)
        
        lao_adm_boundaries[column_name] = None
        for i, row in lao_adm_boundaries.iterrows():
            geom = [row['geometry']]  # Geometry needs to be in a list
            try:
                out_image, out_transform = mask(dataset, geom, crop=True)
                data = out_image[0]  # Read the first band
                lao_adm_boundaries.at[i, column_name] = np.sum(data[data > -99999])
            except ValueError:
                # Handle cases where the geometry does not intersect the raster
                lao_adm_boundaries.at[i, column_name] = 0

# read in the df_merged which contains the population data
import pandas as pd
df_merged = pd.read_csv('lao_pop_projections_with_DNAME.csv')

# merge both dataframes on DCODE, the pop one subset to BTOTL_2020, BTOTL_2024
gdf=lao_adm_boundaries
merged = gdf.merge(df_merged[['DCODE', 'BTOTL_2020', 'BTOTL_2024']], left_on='DCODE', right_on='DCODE')

# create the ratio of BTOTL_2024 to BTOTL_2020, converting each col to integer and removing commas
merged['ratio'] = merged['BTOTL_2024'].str.replace(',', '').astype(int) / merged['BTOTL_2020'].str.replace(',', '').astype(int)

# for each of the lao_X columns, convert them to numeric, and then multiply by the ratio, creating columns lao_X_2024
for column in merged.columns:
    if column.startswith('lao_'):
        if(merged[column].dtype != 'float64'):
            merged[column] = pd.to_numeric(merged[column], errors='coerce')
        merged[column + '_2024'] = merged[column] * merged['ratio']

# Create 2020 and 2024 columns for the sum of all lao_X columns ending in _2020 and _2024
merged['sum_2020'] = merged.filter(regex='lao_.*_2020$').sum(axis=1)
merged['sum_2024'] = merged.filter(regex='lao_.*_2024$').sum(axis=1)

lao_adm_boundaries_union = lao_adm_boundaries[['DCODE', 'TOTAL', 'DNAME','geometry']].dissolve(by='DCODE', aggfunc={'TOTAL':'sum', 'DNAME':'first'})
lao_adm_boundaries_union = lao_adm_boundaries_union.reset_index()

import h3
import geopandas as gpd
from shapely.geometry import Polygon, mapping

# Step 1: Define your bounding box
minx, miny, maxx, maxy = lao_adm_boundaries.total_bounds

# Step 2: Flip coordinates to lat/lon format for h3
bbox_polygon = Polygon([(miny, minx), (miny, maxx), (maxy, maxx), (maxy, minx), (miny, minx)])  # lat/lon format
bbox_geojson = mapping(bbox_polygon)  # Convert to GeoJSON-like structure


# Step 3: Adjust the resolution for larger hexagons initially (e.g., 6)
resolution = 8  # Lower resolution for larger hexagons

# Step 4: Get the list of H3 hexagons that overlap the bounding box
hexagons = h3.polyfill(bbox_geojson, resolution)

# Step 5: Convert H3 hexagons to polygons and store them in a list
hexagon_polygons = []
hex_ids = []
for hex_id in hexagons:
    hex_boundary = h3.h3_to_geo_boundary(hex_id, geo_json=True)
    hexagon_polygons.append(Polygon(hex_boundary))
    hex_ids.append(hex_id)

# Step 6: Create a GeoDataFrame from the list of polygons and hex IDs
hexagons_gdf = gpd.GeoDataFrame({'hex_id': hex_ids, 'geometry': hexagon_polygons}, crs="EPSG:4326")

hexagons_gdf = hexagons_gdf.to_crs(epsg=32648)
lao_adm_boundaries_union = lao_adm_boundaries_union.to_crs(epsg=32648)

# Step 1: Perform a spatial join to find which hexagons intersect with which districts
hex_districts = gpd.sjoin(hexagons_gdf, lao_adm_boundaries_union, how='inner', predicate='intersects')

# Step 2: Merge with the original GeoDataFrame to get the district geometries
hex_districts = hex_districts.merge(lao_adm_boundaries_union[['DCODE', 'geometry']], on='DCODE', suffixes=('', '_district'))

# Step 3: Calculate the intersection area for each hexagon-district pair
hex_districts['intersection_area'] = hex_districts.apply(
    lambda row: row['geometry'].intersection(row['geometry_district']).area,
    axis=1
)

# Step 4: Calculate the fraction of each hexagon within each district
hex_districts['hex_area'] = hex_districts['geometry'].area
hex_districts['fraction'] = hex_districts['intersection_area'] / hex_districts['hex_area']

# merge hex_districts with df_merged on DCODE
hex_districts = hex_districts.merge(df_merged, on='DCODE')

import os
import rasterio
import numpy as np
import geopandas as gpd
from rasterio.mask import mask

filenames = os.listdir('popStructure')[1:]

for filename in filenames:
    if filename.endswith('.tif'):
        column_name = os.path.splitext(filename)[0]
        print("Processing", column_name)
        dataset = rasterio.open('popStructure/' + filename)
        print("pop raster CRS:", dataset.crs)
        gdf[column_name] = None
        for i, row in gdf.iterrows():
            geom = [row['geometry']]  # Geometry needs to be in a list
            out_image, out_transform = mask(dataset, geom, crop=True)
            data = out_image[0]  # Read the first band
            gdf.at[i, column_name] = np.sum(data[data > -99999])

import rpy2.robjects as ro
from rpy2.robjects import pandas2ri

# Activate the automatic conversion between pandas and R data frames
pandas2ri.activate()

# subset hexagons df to all columns except geometry
ro.globalenv['mydf'] = hexagons.drop(columns='geometry')

scaled=ro.conversion.rpy2py(ro.globalenv['scaled'])
# from hexagons, concatenate DCODE, DNAME_x by hex_id and take the first geometry by hex_id
grouped = hexagons.groupby('hex_id').agg({
    'DCODE': concatenate,
    'DNAME_x': concatenate,
    'geometry': 'first'
}).reset_index()
# merge grouped with scaled on hex_id and rename DNAME_x to DNAME

grouped = grouped.merge(scaled, on='hex_id')
# cast to GeoDataFrame using the same parameters as hexagons
grouped = gpd.GeoDataFrame(grouped, crs=hexagons.crs, geometry='geometry')