#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于token分布距离的重采样脚本
读取octuple_token_analysis_report.json，计算每个数据与整体分布的距离，
按照距离加权采样，距离越远的越容易被采样
"""

import os
import numpy as np
from pathlib import Path
from collections import Counter
from tqdm import tqdm
import json
from scipy.stats import entropy, wasserstein_distance
from scipy.spatial.distance import jensenshannon
from concurrent.futures import ThreadPoolExecutor, as_completed
import multiprocessing as mp

# Octuple的列名定义
COLUMN_NAMES = [
    "pitch",      # 0: Pitch/PitchDrum
    "position",   # 1: Position
    "bar",        # 2: Bar
    "velocity",   # 3: Velocity
    "duration",   # 4: Duration
    "program",    # 5: Program
    "tempo",      # 6: Tempo
    "timesig"     # 7: TimeSignature
]


def load_distribution_from_json(json_path):
    """
    从JSON文件中加载整体token分布
    
    Args:
        json_path: JSON文件路径
        
    Returns:
        dict: {column_name: {token: probability}}
    """
    print(f"读取分布文件: {json_path}")
    with open(json_path, 'r', encoding='utf-8') as f:
        report = json.load(f)
    
    distributions = {}
    columns = report.get('columns', {})
    
    for col_name in COLUMN_NAMES:
        if col_name not in columns:
            print(f"警告: 列 {col_name} 不在报告中")
            distributions[col_name] = {}
            continue
        
        col_data = columns[col_name]
        token_counts = col_data.get('token_counts', {})
        total_tokens = col_data.get('total_tokens', 1)
        
        # 转换为概率分布
        distribution = {}
        for token_str, count in token_counts.items():
            token = int(token_str)
            distribution[token] = count / total_tokens
        
        distributions[col_name] = distribution
        print(f"  列 {col_name}: {len(distribution)} 个唯一token, 总token数: {total_tokens:,}")
    
    return distributions


def compute_data_distribution(data, col_idx):
    """
    计算单个数据在指定列的token分布
    
    Args:
        data: numpy数组 (num_tokens, num_columns)
        col_idx: 列索引
        
    Returns:
        dict: {token: probability}
    """
    if data.size == 0:
        return {}
    
    tokens = data[:, col_idx]
    unique, counts = np.unique(tokens, return_counts=True)
    total = len(tokens)
    
    distribution = {}
    for token, count in zip(unique, counts):
        distribution[int(token)] = count / total
    
    return distribution


def compute_emd_distance(dist1, dist2):
    """
    使用推土机距离（Earth Mover's Distance / Wasserstein距离）计算两个分布之间的距离
    
    Args:
        dist1: 分布1，dict {token: probability}，已归一化
        dist2: 分布2，dict {token: probability}，已归一化
        
    Returns:
        float: EMD距离
    """
    # 获取所有token的并集，并排序
    all_tokens = sorted(set(dist1.keys()) | set(dist2.keys()))
    
    if not all_tokens:
        return 0.0
    
    # 构建概率向量和token值向量
    p_weights = np.array([dist1.get(token, 0.0) for token in all_tokens])
    q_weights = np.array([dist2.get(token, 0.0) for token in all_tokens])
    token_values = np.array(all_tokens, dtype=float)
    
    # 归一化（处理数值误差）
    p_sum = p_weights.sum()
    q_sum = q_weights.sum()
    
    if p_sum < 1e-10 or q_sum < 1e-10:
        return 0.0
    
    p_weights = p_weights / p_sum
    q_weights = q_weights / q_sum
    
    # 使用Wasserstein距离（1-Wasserstein距离，即推土机距离）
    # wasserstein_distance需要两个分布的样本值位置和权重
    # 对于离散分布，我们使用token值作为位置
    emd = wasserstein_distance(token_values, token_values, p_weights, q_weights)
    
    return emd


def compute_distribution_distance(dist1, dist2, method='emd'):
    """
    计算两个分布之间的距离
    
    Args:
        dist1: 分布1，dict {token: probability}
        dist2: 分布2，dict {token: probability}
        method: 距离计算方法，'emd' (推土机距离), 'js' (Jensen-Shannon) 或 'kl' (KL散度)
        
    Returns:
        float: 分布距离
    """
    if method == 'emd':
        return compute_emd_distance(dist1, dist2)
    
    # 获取所有token的并集
    all_tokens = set(dist1.keys()) | set(dist2.keys())
    
    if not all_tokens:
        return 0.0
    
    # 构建概率向量
    p = np.array([dist1.get(token, 0.0) for token in all_tokens])
    q = np.array([dist2.get(token, 0.0) for token in all_tokens])
    
    # 归一化（处理数值误差）
    p = p / (p.sum() + 1e-10)
    q = q / (q.sum() + 1e-10)
    
    if method == 'js':
        # Jensen-Shannon散度（对称，范围[0, 1]）
        return jensenshannon(p, q)
    elif method == 'kl':
        # KL散度（非对称，需要处理零值）
        # 添加小的平滑项避免log(0)
        p = p + 1e-10
        q = q + 1e-10
        p = p / p.sum()
        q = q / q.sum()
        return entropy(p, q)
    else:
        raise ValueError(f"未知的距离方法: {method}")


def extract_subdistribution(global_dist, data_tokens):
    """
    从全局分布中提取只包含数据中出现的token的子分布，并归一化
    
    Args:
        global_dist: 全局分布，dict {token: probability}
        data_tokens: 数据中出现的token集合，set或list
        
    Returns:
        dict: 子分布，dict {token: probability}，已归一化
    """
    if not data_tokens or not global_dist:
        return {}
    
    # 提取子分布
    sub_dist = {token: global_dist.get(token, 0.0) for token in data_tokens}
    
    # 归一化
    total = sum(sub_dist.values())
    if total < 1e-10:
        return {}
    
    normalized_sub_dist = {token: prob / total for token, prob in sub_dist.items()}
    
    return normalized_sub_dist


def compute_data_distance(data, global_distributions, method='emd'):
    """
    计算单个数据与整体分布的距离
    对每首歌，从数据集分布中找出和这首歌的分布包含的数据相同的子分布，
    都进行归一化然后计算推土机距离
    
    Args:
        data: numpy数组 (num_tokens, num_columns) 或文件路径（如果是延迟加载）
        global_distributions: 整体分布，dict {column_name: {token: probability}}
        method: 距离计算方法，'emd' (推土机距离), 'js' (Jensen-Shannon) 或 'kl' (KL散度)
        
    Returns:
        float: 平均距离（跨所有列）
    """
    # 如果data是路径，则加载它
    if isinstance(data, (str, Path)):
        try:
            data = np.load(data)['arr_0']
        except Exception as e:
            # 不打印错误，让调用者处理
            raise RuntimeError(f"加载文件 {data} 时出错: {e}")
    
    distances = []
    
    for col_idx, col_name in enumerate(COLUMN_NAMES):
        # 计算该数据在该列的分布
        data_dist = compute_data_distribution(data, col_idx)
        
        # 获取整体分布
        global_dist = global_distributions.get(col_name, {})
        
        if not data_dist or not global_dist:
            continue
        
        # 从全局分布中提取只包含数据中出现的token的子分布
        data_tokens = set(data_dist.keys())
        sub_global_dist = extract_subdistribution(global_dist, data_tokens)
        
        if not sub_global_dist:
            continue
        
        # 归一化数据分布
        data_dist_sum = sum(data_dist.values())
        if data_dist_sum < 1e-10:
            continue
        normalized_data_dist = {token: prob / data_dist_sum 
                               for token, prob in data_dist.items()}
        
        # 计算距离（两个分布都已归一化）
        dist = compute_distribution_distance(normalized_data_dist, sub_global_dist, method=method)
        distances.append(dist)
    
    # 返回平均距离
    return np.mean(distances) if distances else 0.0


def _load_single_file(npz_file):
    """
    加载单个npz文件的辅助函数（用于多线程）
    
    Args:
        npz_file: npz文件路径
        
    Returns:
        tuple: (data, file_path) 或 None（如果加载失败）
    """
    try:
        data = np.load(npz_file)['arr_0']
        if data.ndim == 2:
            return (data, npz_file)
        elif data.ndim == 1:
            print(f"警告: {npz_file} 是一维数组，跳过")
            return None
    except Exception as e:
        print(f"错误: 加载 {npz_file} 时出错: {e}")
        return None


def get_data_file_paths(data_dir):
    """
    获取所有数据文件路径（不加载数据）
    
    Args:
        data_dir: 数据目录路径
        
    Returns:
        list: 文件路径列表
    """
    data_dir = Path(data_dir)
    npz_files = []
    
    if data_dir.exists():
        npz_files = sorted(list(data_dir.rglob("*.npz")))
    
    if not npz_files:
        print(f"警告: 在 {data_dir} 中未找到.npz文件")
        return []
    
    print(f"找到 {len(npz_files)} 个.npz文件")
    return npz_files


def load_data_with_paths(data_dir, num_threads=None, lazy=False):
    """
    加载所有数据并返回数据路径列表（多线程版本）
    
    Args:
        data_dir: 数据目录路径
        num_threads: 线程数，None表示使用CPU核心数
        lazy: 如果为True，只返回文件路径，不加载数据
        
    Returns:
        tuple: (data_list, file_paths_list) 或 (None, file_paths_list) 如果lazy=True
    """
    data_dir = Path(data_dir)
    npz_files = []
    
    if data_dir.exists():
        npz_files = sorted(list(data_dir.rglob("*.npz")))
    
    if not npz_files:
        print(f"警告: 在 {data_dir} 中未找到.npz文件")
        return [], []
    
    if lazy:
        print(f"找到 {len(npz_files)} 个.npz文件（延迟加载模式）")
        return None, npz_files
    
    print(f"找到 {len(npz_files)} 个.npz文件，开始加载...")
    
    if num_threads is None:
        num_threads = min(mp.cpu_count(), len(npz_files))
    
    all_data = []
    file_paths = []
    
    # 使用多线程加载文件
    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        futures = {executor.submit(_load_single_file, npz_file): npz_file 
                   for npz_file in npz_files}
        
        for future in tqdm(as_completed(futures), total=len(futures), desc="加载数据"):
            result = future.result()
            if result is not None:
                data, file_path = result
                all_data.append(data)
                file_paths.append(file_path)
    
    # 保持原始顺序
    if file_paths:
        sorted_pairs = sorted(zip(file_paths, all_data), key=lambda x: str(x[0]))
        file_paths, all_data = zip(*sorted_pairs)
        file_paths = list(file_paths)
        all_data = list(all_data)
    
    return all_data, file_paths


def weighted_resample(file_paths, distances, sample_ratio=0.3, method='js', lazy=True):
    """
    根据距离进行加权重采样
    
    Args:
        file_paths: 文件路径列表
        distances: 距离列表
        sample_ratio: 采样比例
        method: 距离计算方法（用于确定权重方向）
        lazy: 如果为True，返回文件路径而不是数据
        
    Returns:
        tuple: (sampled_data_or_paths, sampled_paths, sampled_indices)
    """
    n_samples = int(len(file_paths) * sample_ratio)
    print(f"\n准备采样 {n_samples} 个数据 (占总数的 {sample_ratio*100:.1f}%)")
    
    # 将距离转换为权重
    # 距离越远，权重越大
    distances = np.array(distances)
    
    # 处理零距离或异常值
    min_dist = np.min(distances[distances > 0]) if np.any(distances > 0) else 1e-10
    distances = np.maximum(distances, min_dist * 0.1)
    
    # 归一化距离到[0, 1]，然后转换为权重
    # 使用指数函数增强距离差异
    normalized_distances = (distances - distances.min()) / (distances.max() - distances.min() + 1e-10)
    weights = np.exp(normalized_distances * 3)  # 指数放大，使距离远的更容易被采样
    
    # 归一化权重
    weights = weights / weights.sum()
    
    # 加权随机采样
    indices = np.arange(len(file_paths))
    sampled_indices = np.random.choice(indices, size=n_samples, replace=False, p=weights)
    
    sampled_paths = [file_paths[i] for i in sampled_indices]
    
    # 如果lazy=True，返回路径；否则加载数据
    if lazy:
        sampled_data = sampled_paths  # 返回路径，延迟加载
    else:
        # 加载采样后的数据
        sampled_data = []
        for path in tqdm(sampled_paths, desc="加载采样数据"):
            try:
                data = np.load(path)['arr_0']
                sampled_data.append(data)
            except Exception as e:
                print(f"错误: 加载 {path} 时出错: {e}")
                sampled_data.append(None)
    
    print(f"采样完成:")
    print(f"  采样数据数量: {len(sampled_paths)}")
    print(f"  平均距离: {distances[sampled_indices].mean():.6f}")
    print(f"  最小距离: {distances[sampled_indices].min():.6f}")
    print(f"  最大距离: {distances[sampled_indices].max():.6f}")
    
    return sampled_data, sampled_paths, sampled_indices


def _save_single_file(args_tuple):
    """
    保存单个文件的辅助函数（用于多线程）
    支持延迟加载：如果data是路径，则从文件加载
    
    Args:
        args_tuple: (data, original_path, output_dir)
        
    Returns:
        tuple: (success, original_path) 或 (False, original_path, error_msg)
    """
    data, original_path, output_dir = args_tuple
    try:
        # 如果data是路径，则加载它
        if isinstance(data, (str, Path)):
            data = np.load(data)['arr_0']
        
        # 保持相对路径结构
        relative_path = original_path.relative_to(original_path.parents[len(original_path.parts) - 3])
        output_path = output_dir / relative_path
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        np.savez_compressed(output_path, data)
        return (True, original_path)
    except Exception as e:
        error_msg = str(e)
        print(f"错误: 保存 {original_path} 时出错: {error_msg}")
        return (False, original_path, error_msg)


def save_sampled_data(sampled_data, sampled_paths, output_dir, num_threads=None):
    """
    保存采样后的数据（多线程版本）
    
    Args:
        sampled_data: 采样后的数据列表
        sampled_paths: 采样后的文件路径列表
        output_dir: 输出目录
        num_threads: 线程数，None表示使用CPU核心数
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"\n保存采样数据到: {output_dir}")
    
    if num_threads is None:
        num_threads = min(mp.cpu_count(), len(sampled_data))
    
    # 准备参数
    save_args = [(data, original_path, output_dir) 
                 for data, original_path in zip(sampled_data, sampled_paths)]
    
    # 使用多线程保存文件
    success_count = 0
    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        # 提交所有任务
        futures = [executor.submit(_save_single_file, args) 
                   for args in save_args]
        
        # 收集结果
        for future in tqdm(as_completed(futures), total=len(futures), desc="保存数据"):
            try:
                result = future.result(timeout=300)  # 设置超时避免卡死
                if isinstance(result, tuple) and len(result) >= 2:
                    success = result[0]
                    if success:
                        success_count += 1
            except Exception as e:
                print(f"错误: 获取保存结果时出错: {e}")
    
    print(f"保存完成，共保存 {success_count}/{len(sampled_data)} 个文件")


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="基于token分布距离的重采样")
    parser.add_argument("--json_path", type=str, 
                       default="octuple_token_analysis_report.json",
                       help="token分析报告JSON文件路径")
    parser.add_argument("--data_dir", type=str,
                       default="dataset/represented_data/tuneidx/tuneidx_msmidi/oct8",
                       help="数据目录路径")
    parser.add_argument("--output_dir", type=str,
                       default="dataset/represented_data/tuneidx/tuneidx_msmidi/oct8_resampled",
                       help="输出目录路径")
    parser.add_argument("--sample_ratio", type=float, default=0.3,
                       help="采样比例 (默认: 0.3)")
    parser.add_argument("--distance_method", type=str, default="emd",
                       choices=["emd", "js", "kl"],
                       help="距离计算方法: 'emd' (推土机距离/EMD), 'js' (Jensen-Shannon) 或 'kl' (KL散度)")
    parser.add_argument("--seed", type=int, default=42,
                       help="随机种子")
    parser.add_argument("--num_threads", type=int, default=1,
                       help="线程数，None表示使用CPU核心数 (默认: None)")
    
    args = parser.parse_args()
    
    # 设置随机种子
    np.random.seed(args.seed)
    
    # 1. 加载整体分布
    global_distributions = load_distribution_from_json(args.json_path)
    
    # 2. 获取所有数据文件路径（延迟加载模式，避免一次性加载所有数据）
    _, file_paths = load_data_with_paths(args.data_dir, lazy=True)
    
    if not file_paths:
        print("错误: 未找到任何数据文件")
        return
    
    print(f"\n共找到 {len(file_paths)} 个数据文件")
    
    # 3. 计算每个数据与整体分布的距离（多线程版本，延迟加载）
    print("\n计算每个数据与整体分布的距离（延迟加载模式）...")
    
    def _compute_distance_wrapper(args_tuple):
        """计算距离的包装函数（用于多线程，支持延迟加载）"""
        idx, file_path, global_dists, method = args_tuple
        try:
            distance = compute_data_distance(file_path, global_dists, method=method)
            return (idx, distance, None)
        except Exception as e:
            return (idx, 0.0, str(e))
    
    if args.num_threads is None:
        num_threads = min(mp.cpu_count(), len(file_paths))
    else:
        num_threads = args.num_threads
    
    # 准备参数（使用文件路径而不是数据，包含索引）
    distance_args = [(i, file_path, global_distributions, args.distance_method) 
                     for i, file_path in enumerate(file_paths)]
    
    # 使用多线程计算距离（按需加载数据）
    # 初始化结果列表，保持顺序
    distances = [0.0] * len(file_paths)
    
    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        # 提交所有任务
        futures = [executor.submit(_compute_distance_wrapper, args) 
                   for args in distance_args]
        
        # 收集结果，使用 tqdm 显示进度
        for future in tqdm(as_completed(futures), total=len(futures), desc="计算距离"):
            try:
                idx, distance, error = future.result(timeout=300)  # 设置超时避免卡死
                distances[idx] = distance
                if error:
                    print(f"警告: 计算距离时出错 (索引 {idx}): {error}")
            except Exception as e:
                print(f"错误: 获取结果时出错: {e}")
                # 如果无法获取结果，保持默认值 0.0
    
    distances = np.array(distances)
    print(f"\n距离统计:")
    print(f"  平均距离: {distances.mean():.6f}")
    print(f"  最小距离: {distances.min():.6f}")
    print(f"  最大距离: {distances.max():.6f}")
    print(f"  标准差: {distances.std():.6f}")
    
    # 4. 根据距离进行加权采样（延迟加载模式）
    sampled_data, sampled_paths, sampled_indices = weighted_resample(
        file_paths, distances, 
        sample_ratio=args.sample_ratio,
        method=args.distance_method,
        lazy=True  # 使用延迟加载，避免重复加载数据
    )
    
    # 5. 保存采样结果（多线程，延迟加载）
    save_sampled_data(sampled_data, sampled_paths, args.output_dir, num_threads=args.num_threads)
    
    # 6. 保存采样索引（可选，用于后续分析）
    indices_file = Path(args.output_dir) / "sampled_indices.npy"
    np.save(indices_file, sampled_indices)
    print(f"\n采样索引已保存到: {indices_file}")
    
    # 保存采样信息
    info = {
        "total_samples": len(file_paths),
        "sampled_samples": len(sampled_data),
        "sample_ratio": args.sample_ratio,
        "distance_method": args.distance_method,
        "distance_stats": {
            "mean": float(distances.mean()),
            "min": float(distances.min()),
            "max": float(distances.max()),
            "std": float(distances.std())
        },
        "sampled_distance_stats": {
            "mean": float(distances[sampled_indices].mean()),
            "min": float(distances[sampled_indices].min()),
            "max": float(distances[sampled_indices].max()),
            "std": float(distances[sampled_indices].std())
        }
    }
    
    info_file = Path(args.output_dir) / "resample_info.json"
    with open(info_file, 'w', encoding='utf-8') as f:
        json.dump(info, f, indent=2, ensure_ascii=False)
    print(f"采样信息已保存到: {info_file}")


if __name__ == "__main__":
    main()

