import os
import sys

import cftime
import iris.util
from iris.cube import CubeList, Cube

"""
THIS IS NO LONGER USED NOW WE ARE PROCESSING THE .NC FILES SUPPLIED BY THE MET OFFICE
DELETE AFTER A "COOLING OFF" PERIOD
"""

def save_cube(cube: iris.cube.Cube,
              sanitised_field_name: str,
              year_string: str,
              month_string: str,
              output_path: str):

    os.makedirs(output_path, exist_ok = True)

    file_name = "{}-no_part-{}_{}.nc".format(sanitised_field_name, year_string, month_string)
    nc_file_path = os.path.join(output_path, file_name)

    print("writing {}".format(nc_file_path))
    iris.save(cube, nc_file_path)

    # nc_zip_file_path = "{}.zip".format(nc_file_path)
    # print("zipping to {}".format(nc_zip_file_path))
    # with ZipFile(nc_zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zip:
    #     zip.write(nc_file_path, basename(nc_file_path))
    #
    # print("deleting {}".format(nc_file_path))
    # os.remove(nc_file_path)


def process_met_office_NAME(name_file_wildcard: str, out_path: str):

    cube_list = iris.load_raw(name_file_wildcard)

    # get all the field names from the cubes as a unique set
    field_names: set = set()

    for cube in cube_list:
        field_names.add(cube.name())

    # write a cube for each calendar month for each field type
    for field_name in field_names:

        field_name: str = field_name
        print(field_name)
        cube: Cube = cube_list.extract(iris.Constraint(name = field_name))

        for month in range(1, 13):

            month_constraint = iris.Constraint(time = lambda cell: cell.point.month == month)

            cubes_for_month: CubeList = cube.extract(month_constraint)

            if len(cubes_for_month) != 0:
                cube_for_month: Cube = cubes_for_month.merge_cube()

                # hack, the 'cell_methods' field is NonType and this gets the save method of iris upset,
                # so setting to an empty Tuple
                cube_for_month.cell_methods = ()

                time_coord = cube_for_month.coord('time')
                sample_date = time_coord.units.num2date(time_coord.points[0])
                year_string = cftime.datetime.strftime(sample_date, "%Y")
                month_string = cftime.datetime.strftime(sample_date, "%m")
                # year_string = cftime.real_datetime.strftime(sample_date, "%Y")
                # month_string = cftime.real_datetime.strftime(sample_date, "%m")

                # get rid of funny characters, for filename
                sanitised_field_name = field_name.replace("(", "")
                sanitised_field_name = sanitised_field_name.replace(")", "")
                sanitised_field_name = sanitised_field_name.replace("/", "_")

                save_cube(cube_for_month,
                          sanitised_field_name,
                          year_string,
                          month_string,
                          out_path)


if __name__ == '__main__':

    """
    example usage:
    /media/scratch/lb584_scratch/projects/ews/input/eth/MET_20200707/WR_EnvSuit_Met_Ethiopia_20200707/*.nc
    /media/scratch/lb584_scratch/projects/ews/output/name_processing/another_dir
    """

    _name_file_wildcard = sys.argv[1]
    _out_path = sys.argv[2]

    process_met_office_NAME(_name_file_wildcard, _out_path)