import json
import os
import warnings
import sys
import time
import argparse
import ray
import numpy as np
import Flexbrdf.hytools as ht
from Flexbrdf.hytools.io.envi import *
from Flexbrdf.hytools.topo import calc_topo_coeffs
from Flexbrdf.hytools.brdf import calc_brdf_coeffs
from Flexbrdf.hytools.glint import set_glint_parameters
from Flexbrdf.hytools.masks import mask_create

warnings.filterwarnings("ignore")
np.seterr(divide='ignore', invalid='ignore')

# Default configuration values
DEFAULT_CONFIG = {
    'file_type': 'envi',
    'num_cpus': 16,
    'bad_bands': [],
    'corrections': [],
    'resample': False,
    'resampler': {'type': 'cubic', 'out_waves': []},
    'export': {
        'coeffs':True,
        'image': True,
        'masks': False,
        'subset_waves': [],
        'output_dir': './output/',
        'suffix': 'BRDF'
    },
    'brdf': {
        'type': 'flex',
        'grouped': True,
        'geometric': 'li_dense_r',
        'volume': 'ross_thick',
        'b/r': 2.5,
        'h/b': 2.0,
        'sample_perc': 0.1,
        'interp_kind': 'linear',
        'calc_mask': [
            ['water', {'band_1': 850, 'band_2': 660,"threshold": 290}],
            ['kernel_finite', {}],
            ['ancillary', {'name': 'sensor_zn', 'min': 0.03490658503988659, 'max': 'inf'}]
        ],
        'apply_mask': [
            ['water', {'band_1': 850, 'band_2': 660,"threshold": 290}]
        ],
        'bin_type': 'dynamic',
        'num_bins': 18,
        'ndvi_bin_min': 0.05,
        'ndvi_bin_max': 1.0,
        'ndvi_perc_min': 10,
        'ndvi_perc_max': 95,
        'solar_zn_type': 'scene'
    },
    # 'topo': {
    #     'type': 'scs+c',
    #     'calc_mask': [
    #         ['ndi', {'band_1': 850, 'band_2': 660, 'min': 0.05, 'max': 1.0}],
    #         ['kernel_finite', {}],
    #         ['ancillary', {'name': 'sensor_zn', 'min': 0.03490658503988659, 'max': 'inf'}]
    #     ],
    #     'apply_mask': [
    #         ['ndi', {'band_1': 850, 'band_2': 660, 'min': 0.05, 'max': 1.0}]
    #     ],
    #     'sample_perc': 0.1,
    #     'subgrouped': False,
    #     'subgroup': {}
    # },
    # 'glint': {
    #     'type': 'hochberg',
    #     'correction_band': 560,
    #     'deep_water_sample': {},
    #     'calc_mask': [
    #         ['ndi', {'band_1': 550, 'band_2': 2150, 'min': -1, 'max': 0}],
    #         ['kernel_finite', {}]
    #     ],
    #     'apply_mask': [
    #         ['ndi', {'band_1': 550, 'band_2': 2150, 'min': -1, 'max': 0}]
    #     ]
    # }
}

def build_config_from_args(args):
    """从命令行参数构建配置字典"""
    # 自动发现输入文件
    input_files = []
    if os.path.isdir(args.input_dir):
        # 支持的文件扩展名
        extensions = ['.tif', '.tiff', '.envi', '.img', '.dat']
        for file in os.listdir(args.input_dir):
            if any(file.lower().endswith(ext) for ext in extensions):
                input_files.append(os.path.join(args.input_dir, file))
    else:
        input_files = [args.input_dir] if os.path.isfile(args.input_dir) else []

    if not input_files:
        raise ValueError(f"No input files found in {args.input_dir}")

    # 自动生成anc_files
    anc_files = {}
    for input_file in input_files:
        base_name = os.path.splitext(os.path.basename(input_file))[0]
        # 根据文件名模式生成对应的angles文件路径
        # 例如: 2025_9_2_3_53_45_202592_35252_0_rad_geo_corrected_reflectance.dat
        # 对应的angles文件: 2025_9_2_3_53_45_202592_35252_0_rad_rgbxyz_geo_registered_angles.bip
        parts = base_name.split('_')
        if len(parts) >= 9:  # 确保有足够的parts
            anc_base = f"{parts[0]}_{parts[1]}_{parts[2]}_{parts[3]}_{parts[4]}_{parts[5]}_{parts[6]}_{parts[7]}_{parts[8]}_rad_rgbxyz_geo_angles_registered"
            # 确保路径分隔符统一使用反斜杠（Windows）
            anc_dir_clean = args.anc_dir.replace('/', '\\')
            anc_path = os.path.join(anc_dir_clean, anc_base + ".bip")
            # 确保最终路径格式正确
            anc_path = os.path.normpath(anc_path)

            if os.path.exists(anc_path):
                # 生成完整的anc_files字典结构
                anc_files[input_file] = {
                    "path_length": [anc_path, 0],
                    "sensor_az": [anc_path, 9],
                    "sensor_zn": [anc_path, 8],
                    "solar_az": [anc_path, 7],
                    "solar_zn": [anc_path, 6],
                    "phase": [anc_path, 0],
                    "slope": [anc_path, 0],
                    "aspect": [anc_path, 0],
                    "cosine_i": [anc_path, 0],
                    "utc_time": [anc_path, 0]
                }
            else:
                # 如果找不到对应的angles文件，使用默认的obs文件
                obs_path = os.path.join(args.anc_dir, base_name + "_obs")
                if os.path.exists(obs_path):
                    anc_files[input_file] = obs_path

    # 构建配置字典
    config_dict = DEFAULT_CONFIG.copy()
    config_dict.update({
        'input_files': input_files,
        'anc_files': anc_files,
        'file_type': 'envi',  # 假设都是ENVI格式
        'num_cpus': args.num_cpus,
        'bad_bands': args.bad_bands if args.bad_bands else [],
        'corrections': args.corrections if args.corrections else [],
        'export': {
            'coeffs': args.export_coeffs,
            'image': args.export_image,
            'masks': args.export_masks,
            'subset_waves': args.subset_waves if args.subset_waves else [],
            'output_dir': args.output_dir,
            'suffix': args.suffix
        }
    })

    # 根据命令行参数更新BRDF配置
    if 'brdf' in args.corrections:
        brdf_config = config_dict['brdf'].copy()
        if hasattr(args, 'brdf_type') and args.brdf_type:
            brdf_config['type'] = args.brdf_type
        if hasattr(args, 'grouped') and args.grouped is not None:
            brdf_config['grouped'] = args.grouped
        if hasattr(args, 'geometric') and args.geometric:
            brdf_config['geometric'] = args.geometric
        if hasattr(args, 'volume') and args.volume:
            brdf_config['volume'] = args.volume
        if hasattr(args, 'num_bins') and args.num_bins:
            brdf_config['num_bins'] = args.num_bins
        config_dict['brdf'] = brdf_config

    # 根据命令行参数更新TOPO配置
    if 'topo' in args.corrections:
        topo_config = config_dict['topo'].copy()
        if hasattr(args, 'topo_type') and args.topo_type:
            topo_config['type'] = args.topo_type
        config_dict['topo'] = topo_config

    # 根据命令行参数更新Glint配置
    if 'glint' in args.corrections:
        glint_config = config_dict['glint'].copy()
        if hasattr(args, 'glint_type') and args.glint_type:
            glint_config['type'] = args.glint_type
        config_dict['glint'] = glint_config

    return config_dict

def main():
    parser = argparse.ArgumentParser(description="High-resolution image correction with automatic configuration")

    # 必需参数
    parser.add_argument('input_dir', help='Input directory containing image files or single image file path')
    parser.add_argument('anc_dir', help='Ancillary directory containing angle/obs files')
    parser.add_argument('--output-dir', required=True, help='Output directory for corrected images')

    # 可选参数 - 基本设置
    parser.add_argument('--num-cpus', type=int, default=16, help='Number of CPUs to use (default: 4)')
    parser.add_argument('--bad-bands', nargs='*', type=int, default=[], help='Bad band indices to exclude')

    # 校正类型
    parser.add_argument('--corrections', nargs='*', choices=['topo', 'brdf', 'glint'],
                       default=['brdf'], help='Correction types to apply (default: brdf)')

    # BRDF参数
    parser.add_argument('--brdf-type', choices=['universal', 'flex'], default='flex',
                       help='BRDF correction type (default: flex)')
    parser.add_argument('--grouped', action='store_true', default=True,
                       help='Group images for BRDF correction (default: True)')
    parser.add_argument('--no-grouped', action='store_false', dest='grouped',
                       help='Do not group images for BRDF correction')
    parser.add_argument('--geometric', default='li_dense_r',
                       choices=['li_sparse', 'li_dense', 'li_dense_r', 'roujean'],
                       help='Geometric kernel type (default: li_dense_r)')
    parser.add_argument('--volume', default='ross_thick',
                       choices=['ross_thin', 'ross_thick', 'hotspot', 'roujean'],
                       help='Volume kernel type (default: ross_thick)')
    parser.add_argument('--num-bins', type=int, default=18,
                       help='Number of NDVI bins for FlexBRDF (default: 18)')

    # TOPO参数
    parser.add_argument('--topo-type', default='scs+c',
                       choices=['scs', 'scs+c', 'c', 'cosine', 'mod_minneart'],
                       help='TOPO correction type (default: scs+c)')

    # Glint参数
    parser.add_argument('--glint-type', default='hochberg',
                       choices=['hochberg', 'gao', 'hedley'],
                       help='Glint correction type (default: hochberg)')

    # 输出选项
    parser.add_argument('--export-coeffs', action='store_true', default=True,
                       help='Export correction coefficients')
    parser.add_argument('--export-image', action='store_true', default=True,
                       help='Export corrected images (default: True)')
    parser.add_argument('--export-masks', action='store_true', default=True,
                       help='Export correction masks (default: True)')
    parser.add_argument('--subset-waves', nargs='*', type=float, default=[],
                       help='Subset of wavelengths to export (empty for all)')
    parser.add_argument('--suffix', default='corrected',
                       help='Suffix for output files (default: corrected)')
    parser.add_argument('--output-config', type=str, default=None,
                       help='Output current configuration to JSON file')

    # 向后兼容：如果只有一个参数且是JSON文件，则使用传统模式
    if len(sys.argv) == 2 and sys.argv[1].endswith('.json'):
        config_file = sys.argv[1]
        with open(config_file, 'r') as outfile:
            config_dict = json.load(outfile)
    else:
        args = parser.parse_args()
        config_dict = build_config_from_args(args)

    # 如果指定了输出配置选项，则输出配置到JSON文件
    if hasattr(args, 'output_config') and args.output_config:
        print(f"Outputting configuration to {args.output_config}...")
        # 重新排序配置字典以匹配期望的格式
        ordered_config = {
            "bad_bands": config_dict.get("bad_bands", []),
            "file_type": config_dict.get("file_type", "envi"),
            "input_files": config_dict.get("input_files", []),
            "anc_files": config_dict.get("anc_files", {}),
            "num_cpus": config_dict.get("num_cpus", 4),
            "export": config_dict.get("export", {}),
            "corrections": config_dict.get("corrections", []),
            "brdf": config_dict.get("brdf", {}),
            "topo": config_dict.get("topo", {}),
            "glint": config_dict.get("glint", {}),
            "resample": config_dict.get("resample", False),
            "resampler": config_dict.get("resampler", {"type": "cubic", "out_waves": []})
        }
        with open(args.output_config, 'w', encoding='utf-8') as f:
            json.dump(ordered_config, f, indent=2, ensure_ascii=False)
        print("Configuration output complete.")
        return  # 如果只是输出配置，则退出程序

    print("Starting image correction process...")
    total_start = time.time()

    images = config_dict["input_files"]

    if ray.is_initialized():
        ray.shutdown()
    print("Using %s CPUs." % config_dict['num_cpus'])
    ray.init(num_cpus = config_dict['num_cpus'])

    HyTools = ray.remote(ht.HyTools)
    actors = [HyTools.remote() for image in images]

    print(f"Loading {len(images)} image(s)...")
    load_start = time.time()

    if config_dict['file_type'] == 'envi':
        anc_files = config_dict["anc_files"]
        _ = ray.get([a.read_file.remote(image,config_dict['file_type'],
                                        anc_files[image]) for a,image in zip(actors,images)])

    elif config_dict['file_type'] == 'neon':
        _ = ray.get([a.read_file.remote(image,config_dict['file_type']) for a,image in zip(actors,images)])

    _ = ray.get([a.create_bad_bands.remote(config_dict['bad_bands']) for a in actors])

    load_time = time.time() - load_start
    print(".2f")

    # Apply corrections
    if config_dict["corrections"]:
        print(f"Applying {len(config_dict['corrections'])} correction(s): {', '.join(config_dict['corrections'])}")
        correction_start = time.time()

        for correction in config_dict["corrections"]:
            if correction =='topo':
                calc_topo_coeffs(actors,config_dict['topo'])
            elif correction == 'brdf':
                calc_brdf_coeffs(actors,config_dict)
            elif correction == 'glint':
                set_glint_parameters(actors,config_dict)

        correction_time = time.time() - correction_start
        print(".2f")

    if config_dict['export']['coeffs'] and len(config_dict["corrections"]) > 0:
        print("Exporting correction coefficients...")
        coeff_start = time.time()
        _ = ray.get([a.do.remote(export_coeffs,config_dict['export']) for a in actors])
        coeff_time = time.time() - coeff_start
        print(".2f")

    if config_dict['export']['image']:
        print("Exporting corrected images...")
        export_start = time.time()
        _ = ray.get([a.do.remote(apply_corrections,config_dict) for a in actors])
        export_time = time.time() - export_start
        print(".2f")

    total_time = time.time() - total_start
    print(".2f")

    ray.shutdown()

def export_coeffs(hy_obj,export_dict):
    '''Export correction coefficients to file.
    '''
    for correction in hy_obj.corrections:
        coeff_file = export_dict['output_dir']
        coeff_file += os.path.splitext(os.path.basename(hy_obj.file_name))[0]
        coeff_file += "_%s_coeffs_%s.json" % (correction,export_dict["suffix"])

        with open(coeff_file, 'w') as outfile:
            if correction == 'topo':
                corr_dict = hy_obj.topo
            elif correction == 'glint':
                continue
            else:
                corr_dict = hy_obj.brdf
            json.dump(corr_dict,outfile)

def apply_corrections(hy_obj,config_dict):
    '''Apply correction to image and export
        to file.
    '''

    header_dict = hy_obj.get_header()
    header_dict['data ignore value'] = hy_obj.no_data
    header_dict['data type'] = 12

    output_name = config_dict['export']['output_dir']
    output_name += os.path.splitext(os.path.basename(hy_obj.file_name))[0]
    output_name +=  "_%s" % config_dict['export']["suffix"]

    #Export all wavelengths
    if len(config_dict['export']['subset_waves']) == 0:

        if config_dict["resample"] == True:
            hy_obj.resampler = config_dict['resampler']
            waves= hy_obj.resampler['out_waves']
        else:
            waves = hy_obj.wavelengths

        header_dict['bands'] = len(waves)
        header_dict['wavelength'] = waves

        writer = WriteENVI(output_name,header_dict)
        iterator = hy_obj.iterate(by='line', corrections=hy_obj.corrections,
                                  resample=config_dict['resample'])
        while not iterator.complete:
            line = iterator.read_next()
            writer.write_line(line,iterator.current_line)
        writer.close()

    #Export subset of wavelengths
    else:
        waves = config_dict['export']['subset_waves']
        bands = [hy_obj.wave_to_band(x) for x in waves]
        waves = [round(hy_obj.wavelengths[x],2) for x in bands]
        header_dict['bands'] = len(bands)
        header_dict['wavelength'] = waves

        writer = WriteENVI(output_name,header_dict)
        for b,band_num in enumerate(bands):
            band = hy_obj.get_band(band_num,
                                   corrections=hy_obj.corrections)
            writer.write_band(band, b)
        writer.close()

    #Export masks
    if (config_dict['export']['masks']) and (len(config_dict["corrections"]) > 0):
        masks = []
        mask_names = []

        for correction in config_dict["corrections"]:
            for mask_type in config_dict[correction]['apply_mask']:
                mask_names.append(correction + '_' + mask_type[0])
                masks.append(mask_create(hy_obj, [mask_type]))

        header_dict['data type'] = 1
        header_dict['bands'] = len(masks)
        header_dict['band names'] = mask_names
        header_dict['samples'] = hy_obj.columns
        header_dict['lines'] = hy_obj.lines
        header_dict['wavelength'] = []
        header_dict['fwhm'] = []
        header_dict['wavelength units'] = ''
        header_dict['data ignore value'] = 255


        output_name = config_dict['export']['output_dir']
        output_name += os.path.splitext(os.path.basename(hy_obj.file_name))[0]
        output_name +=  "_%s_mask" % config_dict['export']["suffix"]

        writer = WriteENVI(output_name,header_dict)

        for band_num,mask in enumerate(masks):
            mask = mask.astype(int)
            mask[~hy_obj.mask['no_data']] = 255
            writer.write_band(mask,band_num)

        del masks

if __name__== "__main__":
    main()
