from shapely.affinity import affine_transform, translate
import geopandas as gpd
import pandas as pd

df = pd.read_parquet("data/mitma_v1/zonificacion/mitma_zones.parquet")
df = df[['cpro', 'provincia']]
gdf = gpd.read_file("data/mitma_v1/zonificacion/distritos_mitma.dbf")
gdf = gdf.set_index("ID")

gdf = pd.merge(gdf, df, left_index=True, right_index=True)

xoff=500_000
yoff=500_000

mask = (gdf["cpro"] == "38") | (gdf["cpro"] == "35")
gdf.loc[mask, "geometry"] = gdf.loc[mask, "geometry"].apply(
    lambda geom: translate(geom, xoff=xoff, yoff=yoff)
)

gdf.to_parquet("data/mitma_v1/zonificacion.parquet")

df_pop = pd.read_csv("data/covid19/province_population.csv", index_col="pro_id", dtype={"pro_id": str}  )

gdf_pro = gpd.read_file("data/covid19/provincias.geojson").set_index("id")
gdf_pro = pd.merge(gdf_pro, df_pop, left_index=True, right_index=True)

xoff = 5
yoff = 6.5
matrix = [1, 0, 0, 1, xoff, yoff]

islands = ["38", "35"]
gdf_pro.loc[islands, "geometry"] = gdf_pro.loc[islands, "geometry"].apply(
    lambda geom: affine_transform(geom, matrix)
)

gdf_pro.plot(column="total", cmap="OrRd", legend=True, figsize=(10, 10))
gdf_pro.to_parquet("data/covid19/provincias.parquet")