import datetime
import io
import os
import requests

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import streamlit as st
import xarray as xr

from geogif import gif
from rasterio import features

provinces = gpd.read_file("./data/prov_stat.geojson")
geoms = provinces.geometry.boundary

image = features.rasterize(
    geoms,
    out_shape=[381, 1081],
    transform=(0.10000000149011612, 0.0, -160.0, 0.0, 0.10000000149011612, 32.0),
)

cmap = st.selectbox("Colormap", plt.colormaps(), index=1)

fps = st.slider("Frames per second", 0.01, 50.0, 16.0)

show_dt = st.checkbox("Show timestamp", True)

fetch_data = st.checkbox("Refresh data", False)

if st.button("Create GIF"):
    smoke_url = "https://firesmoke.ca/forecasts/current/dispersion.nc"
    smoke_nc = "./data/dispersion.nc"
    if not os.path.isfile(smoke_nc) or fetch_data:
        response = requests.head(smoke_url)

        total_size = int(response.headers.get("Content-Length"))

        progress_text = "Downloading data. Please wait."
        progress_bar = st.progress(0, text=progress_text)

        size_downloaded = 0
        chunk_size = 8192

        with requests.get(smoke_url, stream=True) as r:
            r.raise_for_status()
            with open(smoke_nc, "wb") as f:
                for chunk in r.iter_content(chunk_size=chunk_size):
                    f.write(chunk)

                    size_downloaded += chunk_size
                    percent_complete = min(size_downloaded / total_size, 1)
                    progress_bar.progress(percent_complete, text=progress_text)

        progress_bar.empty()

    ds_disk = xr.open_dataset(smoke_nc)

    tstep = ds_disk.TFLAG

    dts = []
    for i in tstep:
        data = i.data[0]
        dt = datetime.datetime.strptime(
            str(data[0]) + str(data[1]).zfill(6), "%Y%j%H%M%S"
        )
        dts.append(dt)

    pm25 = ds_disk.PM25

    pm25["TSTEP"] = dts

    for idx, i in enumerate(pm25.data):
        pm25[idx].data[0] = np.maximum(i[0], image * 6)

    buffer = io.BytesIO()

    if show_dt:
        date_format = "%Y-%m-%d %H:%M:%S"
    else:
        date_format = None

    gif(
        pm25[:, :, ::-1, :],
        date_format=date_format,
        fps=fps,
        cmap=cmap,
        date_color=(255, 255, 255),
        to=buffer,
    )

    first_dt = dts[0]
    last_dt = dts[-1]

    st.image(buffer, caption=f"{first_dt} to {last_dt}")
    st.download_button(
        label="Download GIF", data=buffer, file_name="firesmoke.gif", mime="image/gif"
    )

    with st.expander("Data Sources"):
        st.markdown(f"Smoke forecast: [FireSmoke Canada](https://firesmoke.ca/)")
        st.markdown(
            f"Administrative boundaries: [Natural Earth, 1:50M Admin 1 – States, provinces](https://www.naturalearthdata.com/http//www.naturalearthdata.com/download/50m/cultural/ne_50m_admin_1_states_provinces.zip)"
        )