from datetime import datetime, timedelta
import glob
import os

from fire import Fire
import h5py
from matplotlib import pyplot as plt
import numpy as np
import pyart
import pathlib
import cartopy
import wradlib as wrl
import dask.array as da
import pandas as pd
import xarray as xr
import re

from ldcast import forecast
from ldcast.visualization import plots


MAP_PROJECTION = cartopy.crs.PlateCarree()
#2020-12-31T17:10
def convert_to_datetime(filename=""):
    year = int(filename[0:4])
    month = int(filename[5:7])
    day = int(filename[8:10])
    hour = int(filename[11:13])
    minute = int(filename[14:16])
    return datetime(year,month,day,hour,minute)

def filter_files_by_format(directory, pattern):
    filtered_files = []
    regex = re.compile(pattern)
    for filename in os.listdir(directory):
        if regex.match(filename):
            #if filename.endswith("00"):
            #    filtered_files.append(filename)
            if filename.endswith("0"):
                filtered_files.append(filename)
    return filtered_files

def read_multiple_data(
    data_dir,
    path_file
):  
    path_file = os.path.join(data_dir, path_file)
    radar = pyart.io.read(path_file)
    processed_grid = pyart.map.grid_from_radars(
        radar,
        # grid_shape=(1, 1000, 1000),
        grid_shape=(20, 640, 640),
        grid_limits=((1000, 5000), (-300_000, 300_000), (-300_000, 300_000)),
    )
    reflectivity_data = processed_grid.fields['reflectivity']['data'].data
    cmax_data = np.array(np.max(reflectivity_data, axis=0))
    radar_Z = wrl.trafo.idecibel(cmax_data)
    R = wrl.zr.z_to_r(radar_Z, a=200.0, b=1.6)
    processed_grid.fields['reflectivity']['data'].data[:] = R
    R = processed_grid.fields['reflectivity']['data'].data[0]
    # extra_left, extra_right = 128, 128
    # extra_top, extra_bottom = 32, 32

    # R = np.pad(R, ((extra_top, extra_bottom), (extra_left, extra_right)),
    #     mode='constant', constant_values=0) 

    return R



def demo(
    out_dir="/data/data_WF/ldcast_precipitation/reflectivity/train_GT",
    data_dir="/data/data_WF/NhaBe/2022",
    t0=datetime(2023,6,25,0,0),
    interval=timedelta(minutes=10),
    past_timesteps=4,
):
    temp = [ 'T02-2022Pro-Raw (01-19)', 'T02-2022Pro-Raw (20-28)', 'T03-2022Pro-Raw (20-31)', 'T04-2022Pro-Raw (01-16)-',
'T04-2022Pro-Raw (17-30)-', 'T05-2022Pro-Raw (01-11)', 'T05-2022Pro-Raw (12-22)', 'T05-2022Pro-Raw (23-31)' , 'T06- 2022Pro-Raw (01-13)', 'T06- 2022Pro-Raw (14-23)',  'T06- 2022Pro-Raw (24-30)']
    # for subdir in os.listdir(directory_path):
    for subdir in temp:
    # for subdir in os.listdir(data_dir):
        subdir_path = os.path.join(data_dir, subdir)
        if os.path.isdir(subdir_path):
            temp2 = os.listdir(subdir_path)
            temp2.sort()
            for subsubdir in temp2: 
                subsubdir_path = os.path.join(subdir_path, subsubdir)
                pattern_to_filter = r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}" 
                filtered_files = filter_files_by_format(subsubdir_path, pattern_to_filter)
                sorted_files = sorted(filtered_files)
                print("Filtered Files:", len(filtered_files))  
                for i in range(0, len(sorted_files)):
                    t = convert_to_datetime(sorted_files[i])
                    print(t)
                    R_gt = read_multiple_data(subsubdir_path, sorted_files[i])
                    fn = os.path.join(out_dir, f"{t:%Y%m%d%H%M}.npy")
                    np.save(fn, R_gt)


if __name__ == "__main__":
    Fire(demo)