import argparse
import json
import os
import ray
import hytools as ht
from hytools.io.envi import WriteENVI
from hytools.brdf import calc_brdf_coeffs


def build_anc_mapping(anc_file, anc_names):
    return dict(zip(anc_names, [[anc_file, i] for i in range(len(anc_names))]))


def build_brdf_config(args):
    mask = [["ndi", {"band_1": args.ndi_band_1, "band_2": args.ndi_band_2, "min": args.ndi_min, "max": args.ndi_max}]]
    brdf = {
        "type": args.brdf_type,
        "grouped": False,
        "geometric": args.geometric,
        "volume": args.volume,
        "b/r": args.b_r,
        "h/b": args.h_b,
        "sample_perc": args.sample_perc,
        "calc_mask": mask,
        "apply_mask": mask,
        "solar_zn_type": args.solar_zn_type
    }
    if args.brdf_type == "universal":
        brdf["diagnostic_plots"] = False
        brdf["diagnostic_waves"] = []
    if args.brdf_type == "flex":
        brdf["interp_kind"] = "linear"
        brdf["bin_type"] = "dynamic"
        brdf["num_bins"] = args.num_bins
        brdf["ndvi_bin_min"] = args.ndvi_bin_min
        brdf["ndvi_bin_max"] = args.ndvi_bin_max
        brdf["ndvi_perc_min"] = args.ndvi_perc_min
        brdf["ndvi_perc_max"] = args.ndvi_perc_max
    return brdf


def export_brdf_corrected(hy_obj, args):
    header_dict = hy_obj.get_header()
    header_dict["data ignore value"] = hy_obj.no_data
    header_dict["data type"] = 4
    output_name = os.path.join(
        args["output_dir"],
        f"{os.path.splitext(os.path.basename(hy_obj.file_name))[0]}_{args['suffix']}"
    )
    writer = WriteENVI(output_name, header_dict)
    iterator = hy_obj.iterate(by="line", corrections=hy_obj.corrections)
    while not iterator.complete:
        line = iterator.read_next()
        writer.write_line(line, iterator.current_line)
    writer.close()


def export_brdf_coeffs(hy_obj, args):
    output_name = os.path.join(
        args["output_dir"],
        f"{os.path.splitext(os.path.basename(hy_obj.file_name))[0]}_brdf_coeffs_{args['suffix']}.json"
    )
    with open(output_name, "w") as outfile:
        json.dump(hy_obj.brdf, outfile)


def main():
    parser = argparse.ArgumentParser(description="BRDF correction for ENVI BIP images")
    parser.add_argument("image", type=str)
    parser.add_argument("anc_file", type=str)
    parser.add_argument("output_dir", type=str)
    parser.add_argument("--anc-map", type=str, default="")
    parser.add_argument("--anc-names", type=str, default="path_length,sensor_az,sensor_zn,solar_az,solar_zn,phase,slope,aspect,cosine_i,utc_time")
    parser.add_argument("--brdf-type", type=str, default="universal", choices=["universal", "flex"])
    parser.add_argument("--suffix", type=str, default="brdf")
    parser.add_argument("--num-cpus", type=int, default=1)
    parser.add_argument("--bad-bands-json", type=str, default="")
    parser.add_argument("--solar-zn-type", type=str, default="scene")
    parser.add_argument("--geometric", type=str, default="li_dense_r")
    parser.add_argument("--volume", type=str, default="ross_thick")
    parser.add_argument("--b-r", type=float, default=2.5)
    parser.add_argument("--h-b", type=float, default=2.0)
    parser.add_argument("--sample-perc", type=float, default=0.1)
    parser.add_argument("--ndi-band-1", type=float, default=850.0)
    parser.add_argument("--ndi-band-2", type=float, default=660.0)
    parser.add_argument("--ndi-min", type=float, default=0.05)
    parser.add_argument("--ndi-max", type=float, default=1.0)
    parser.add_argument("--num-bins", type=int, default=18)
    parser.add_argument("--ndvi-bin-min", type=float, default=0.05)
    parser.add_argument("--ndvi-bin-max", type=float, default=1.0)
    parser.add_argument("--ndvi-perc-min", type=float, default=10.0)
    parser.add_argument("--ndvi-perc-max", type=float, default=95.0)
    parser.add_argument("--export-coeffs", action="store_true")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    if args.anc_map:
        with open(args.anc_map, "r") as infile:
            anc_map = json.load(infile)
    else:
        anc_names = [x.strip() for x in args.anc_names.split(",") if x.strip()]
        anc_map = build_anc_mapping(args.anc_file, anc_names)

    config_dict = {
        "file_type": "envi",
        "input_files": [args.image],
        "anc_files": {args.image: anc_map},
        "corrections": ["brdf"],
        "brdf": build_brdf_config(args),
        "num_cpus": args.num_cpus
    }

    if args.bad_bands_json:
        config_dict["bad_bands"] = json.loads(args.bad_bands_json)

    if ray.is_initialized():
        ray.shutdown()
    ray.init(num_cpus=config_dict["num_cpus"])

    HyTools = ray.remote(ht.HyTools)
    actor = HyTools.remote()
    ray.get(actor.read_file.remote(args.image, "envi", anc_map))
    if "bad_bands" in config_dict:
        ray.get(actor.create_bad_bands.remote(config_dict["bad_bands"]))

    calc_brdf_coeffs([actor], config_dict)

    ray.get(actor.do.remote(export_brdf_corrected, {"output_dir": args.output_dir, "suffix": args.suffix}))
    if args.export_coeffs:
        ray.get(actor.do.remote(export_brdf_coeffs, {"output_dir": args.output_dir, "suffix": args.suffix}))

    ray.shutdown()


if __name__ == "__main__":
    main()