import numpy as np
import os
import struct

def read_envi_header(header_path):
    """
    读取 ENVI 头文件，返回元数据字典。
    """
    metadata = {}
    with open(header_path, 'r') as f:
        lines = f.readlines()
    for line in lines:
        line = line.strip()
        if not line or line.startswith(';'):
            continue
        if '=' in line:
            key, value = line.split('=', 1)
            key = key.strip().lower()
            value = value.strip().strip('{}')
            # 尝试转换数值
            if value.isdigit():
                metadata[key] = int(value)
            elif value.replace('.', '', 1).isdigit():
                metadata[key] = float(value)
            else:
                metadata[key] = value
    return metadata

def write_envi_header(output_header, metadata):
    """
    将元数据写入 ENVI 头文件。
    """
    with open(output_header, 'w') as f:
        f.write('ENVI\n')
        for key, val in metadata.items():
            if isinstance(val, (int, float)):
                f.write(f'{key} = {val}\n')
            else:
                f.write(f'{key} = {{{val}}}\n')

def convert_bsq(input_bsq, output_bsq):
    """
    主转换函数：逐波段读取，缩放并转换数据类型。
    """
    # 获取头文件路径
    input_hdr = input_bsq + '.hdr' if not input_bsq.endswith('.hdr') else input_bsq.replace('.hdr', '')
    if not os.path.exists(input_hdr):
        base, ext = os.path.splitext(input_bsq)
        input_hdr = base + '.hdr'
        if not os.path.exists(input_hdr):
            raise FileNotFoundError(f"头文件 {input_hdr} 不存在")

    # 读取头文件信息
    meta = read_envi_header(input_hdr)
    samples = meta.get('samples')
    lines = meta.get('lines')
    bands = meta.get('bands')
    data_type = meta.get('data type')
    interleave = meta.get('interleave', 'bsq').lower()

    if interleave != 'bsq':
        raise ValueError("本脚本仅支持 BSQ 格式，当前 interleave = {}".format(interleave))

    # 根据 data_type 确定输入数据类型
    dtype_map = {1: np.uint8, 2: np.int16, 3: np.int32, 4: np.float32, 5: np.float64, 12: np.uint16}
    if data_type not in dtype_map:
        raise ValueError(f"不支持的 data type: {data_type}")
    in_dtype = dtype_map[data_type]

    print(f"输入文件尺寸: 波段数={bands}, 行数={lines}, 列数={samples}")
    print(f"原始 data type = {data_type} ({in_dtype}), 将缩放为 uint16 (type 12)")

    # 计算每个波段的字节数
    band_size = lines * samples * np.dtype(in_dtype).itemsize

    # 打开输入文件，输出文件
    with open(input_bsq, 'rb') as fin, open(output_bsq, 'wb') as fout:
        for band_idx in range(bands):
            # 定位到当前波段的起始位置
            fin.seek(band_idx * band_size)
            # 读取整个波段
            band_data = np.fromfile(fin, dtype=in_dtype, count=lines * samples)
            # 重塑为 (lines, samples)
            band_data = band_data.reshape((lines, samples))

            # 缩放并转换为 uint16
            scaled = np.round(band_data * 10000).astype(np.uint16)

            # 写入输出文件（BSQ 顺序）
            fout.write(scaled.tobytes())

            # 进度提示
            if (band_idx + 1) % 10 == 0 or band_idx == bands - 1:
                print(f"已处理波段 {band_idx+1}/{bands}")

    # 生成新的头文件
    output_hdr = output_bsq + '.hdr'
    new_meta = meta.copy()
    new_meta['data type'] = 12          # uint16
    new_meta['interleave'] = 'bsq'
    # 可选：修改描述
    if 'description' in new_meta:
        new_meta['description'] = new_meta['description'] + ' (scaled by 10000)'
    else:
        new_meta['description'] = 'Scaled to uint16 (factor 10000)'
    write_envi_header(output_hdr, new_meta)

    print(f"转换完成：{output_bsq} 已生成，data type = 12 (uint16)")

if __name__ == '__main__':
    import sys
    if len(sys.argv) != 3:
        print("用法: python script.py input.bsq output.bsq")
        sys.exit(1)
    convert_bsq(sys.argv[1], sys.argv[2])