#!/usr/bin/env python
# coding: utf-8

# In[1]:


# processes Gridded Livestock of the World data


# In[1]:


import os
import sys
import math
import warnings
import yaml
import pandas as pd
import geopandas as gpd
from osgeo import gdal
from pathlib import Path
import glob
import numpy as np
import rasterio.mask
import rasterio
import fiona
from fiona.crs import from_epsg
from pathlib import Path
import folium
# from folium import plugins
# from matplotlib import cm
# from rasterio import warp
# from rasterio.enums import Resampling
from rasterio.features import shapes
from branca.colormap import linear


# In[2]:


city = 'Luanda'


# ### Data processing

# In[3]:


data_folder = Path(r'D:\World Bank\CRP\data\Gridded Livestock of the World')
output_folder = Path('output')
livestocks = {'Bf': 'buffaloes',
              'Ct': 'cattle',
              'Ch': 'chickens',
              'Dk': 'ducks',
              'Gt': 'goats',
              'Ho': 'horses',
              'Pg': 'pigs',
              'Sh': 'sheep'}
clipped_ls = {}


# In[4]:


# read AOI
aoi = gpd.read_file('AOI/luanda.shp').to_crs(epsg = 4326)


# In[49]:


for ls in livestocks:
    input_raster = data_folder / f'{livestocks[ls]}' / f'5_{ls}_2010_Da.tif'
    with rasterio.open(input_raster) as src:
        clipped_ls[ls], out_transform = rasterio.mask.mask(
            src, aoi.geometry, crop = True, all_touched = True)
        out_meta = src.meta.copy()

    out_meta.update({"driver": "GTiff",
                     "height": clipped_ls[ls].shape[1],
                     "width": clipped_ls[ls].shape[2],
                     "transform": out_transform})


# In[50]:


for ls in clipped_ls:
    clipped_ls[ls][clipped_ls[ls] == -1.7e+308] = 0

out_image1 = sum(clipped_ls.values())


# In[51]:


with rasterio.open(output_folder / f'{city.lower()}_livestock.tif', 'w', **out_meta) as dest:
    dest.write(out_image1)


# ### Data conversion to JSON

# In[37]:


with rasterio.open(output_folder / f'{city.lower()}_livestock.tif') as src:
    image = np.float32(src.read(1))
    results = ({'properties': {'raster_val': np.float32(v)}, 'geometry': s} for i, (s, v) in enumerate(shapes(image, mask = None, transform = src.transform)))
    geoms = list(results)
    gpd_polygonized_raster = gpd.GeoDataFrame.from_features(geoms)
    gpd_polygonized_raster = gpd_polygonized_raster[gpd_polygonized_raster.raster_val != 0]

gpd_polygonized_raster.to_file(output_folder / f'{city.lower()}_livestock.geojson', driver = 'GeoJSON', crs = 'EPSG:4326')
# gpd_polygonized_raster = gpd_polygonized_raster.to_crs(epsg = 4326)


# ### Map with JSON

# In[5]:


glw = gpd.read_file(output_folder / f'{city.lower()}_livestock.geojson').to_crs('epsg:4326')


# In[6]:


# Create a folium map centered on the AOI
mymap = folium.Map(control_scale = True)
aoi_bounds = aoi.total_bounds.tolist()
mymap.fit_bounds([aoi_bounds[:2][::-1], aoi_bounds[2:][::-1]])

folium.GeoJson(
    aoi,
    style_function=lambda feature: {
        'fillColor': 'transparent',
        'color': 'gray',
        'weight': 3
    },
    control = False).add_to(mymap)


# In[7]:


# Create a colormap based on the 'YlOrRd' scale and the 'raster_val' column
colormap = linear.YlOrRd_09.scale(glw['raster_val'].min(), glw['raster_val'].max())

# Add GeoJson layer with choropleth style
folium.GeoJson(
    glw,
    name='Number of livestocks',
    style_function=lambda feature: {
        'fillColor': colormap(feature['properties']['raster_val']),
        'fillOpacity': 0.7,
        'weight': 0,
    },
    highlight_function=lambda x: {'fillOpacity': 0.9},
    tooltip=folium.GeoJsonTooltip(fields=['raster_val'], aliases = ['Number of livestocks'], labels=True, sticky=False),
).add_to(mymap)

# Add Layer Control to toggle the choropleth layer
folium.LayerControl().add_to(mymap)


# In[8]:


mymap


# In[90]:


mymap.save(html_folder / "glw.html")


# ### Data conversion to PNG (deprecated)

# Convert to PNG to enable folium visualization

# In[28]:


with rasterio.open(output_folder / f'{city.lower()}_livestock.tif') as src:
    Z = src.read(1)
    min_val = np.min(Z)
    max_val = np.max(Z)
    scaled_array = 255 * (Z - min_val) / (max_val - min_val)
    uint8_array = scaled_array.round().astype(np.uint8)
    
    meta = src.meta
    meta['driver'] = 'PNG'
    meta['dtype'] = 'uint8'
    
    bounds = src.bounds   # used in folium visualization later

with rasterio.open(output_folder / f'{city.lower()}_livestock.png', 'w', **meta) as dst:
    dst.write(uint8_array, 1)
    # dst.write_colormap(
    #     1, {
    #         np.min(uint8_array): (255, 245, 235, 255),
    #         np.max(uint8_array): (127,39,4, 255) })


# ### Map with PNG (deprecated)

# In[93]:


# Create a folium map centered on the AOI
mymap = folium.Map(control_scale = True)
aoi_bounds = aoi.total_bounds.tolist()
mymap.fit_bounds([aoi_bounds[:2][::-1], aoi_bounds[2:][::-1]])

folium.GeoJson(
    aoi,
    style_function=lambda feature: {
        'fillColor': 'transparent',
        'color': 'gray',
        'weight': 3
    },).add_to(mymap)


# In[94]:


# Add the raster image as an overlay
folium.raster_layers.ImageOverlay(
    image=uint8_array,
    bounds=[(bounds[1], bounds[0]), (bounds[3], bounds[2])],
    colormap=lambda x: (255,0,0,x),
    # colormap=cm.viridis,
    opacity=0.8,
).add_to(mymap)


# In[95]:


# Add legend manually
legend_html = f"""
    <div style="
        position: fixed;
        bottom: 50px;
        left: 50px;
        width: 200px;
        height: 100px;
        background-color: white;
        border: 2px solid black;
        z-index: 1000;
        padding: 10px;
        ">
        <h4>Livestock Counts</h4>
        <div style="background: linear-gradient(to right, white, red); height: 20px;"></div>
        <div style="display: flex; justify-content: space-between; padding-top: 5px;">
            <span>0</span>
            <span>{"{:,.0f}".format(max_val)}</span>
        </div>
    </div>
"""

mymap.get_root().html.add_child(folium.Element(legend_html))


# In[96]:


mymap