#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于position和duration token的重采样脚本V2
对于每首歌：
1. 如果包含的position和duration不在总数据集前3个，则必定采样
2. 对于包含的，以某个固定的百分比采样
3. 两个条件满足一个即可
"""

import os
import numpy as np
from pathlib import Path
from collections import Counter
from tqdm import tqdm
import json
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_top_tokens_from_json(json_path, column_name, top_k=3):
    """
    从JSON文件中加载指定列的前top_k个最常见的token
    
    Args:
        json_path: JSON文件路径
        column_name: 列名（如'position'或'duration'）
        top_k: 返回前k个最常见的token
        
    Returns:
        set: 前top_k个最常见的token集合
    """
    print(f"读取分布文件: {json_path}")
    with open(json_path, 'r', encoding='utf-8') as f:
        report = json.load(f)
    
    columns = report.get('columns', {})
    
    if column_name not in columns:
        print(f"警告: 列 {column_name} 不在报告中")
        return set()
    
    col_data = columns[column_name]
    token_counts = col_data.get('token_counts', {})
    
    # 按出现次数排序，获取前top_k个
    sorted_tokens = sorted(token_counts.items(), key=lambda x: int(x[1]), reverse=True)
    top_tokens = {int(token_str) for token_str, _ in sorted_tokens[:top_k]}
    
    print(f"  列 {column_name} 的前{top_k}个最常见token: {sorted(top_tokens)}")
    
    return top_tokens


def get_data_tokens(data, col_idx):
    """
    获取单个数据在指定列的所有唯一token
    
    Args:
        data: numpy数组 (num_tokens, num_columns) 或文件路径
        col_idx: 列索引
        
    Returns:
        set: 唯一token集合
    """
    # 如果data是路径，则加载它
    if isinstance(data, (str, Path)):
        try:
            data = np.load(data)['arr_0']
        except Exception as e:
            raise RuntimeError(f"加载文件 {data} 时出错: {e}")
    
    if data.size == 0:
        return set()
    
    tokens = data[:, col_idx]
    unique_tokens = set(int(token) for token in np.unique(tokens))
    
    return unique_tokens


def should_sample_song(data, top_position_tokens, top_duration_tokens, 
                       contain_sample_ratio=0.3, not_contain_sample_ratio=0.9, rng=None):
    """
    判断一首歌是否应该被采样
    
    Args:
        data: numpy数组 (num_tokens, num_columns) 或文件路径
        top_position_tokens: position列的前3个最常见token集合
        top_duration_tokens: duration列的前3个最常见token集合
        contain_sample_ratio: 对于包含前3个token的歌曲，采样比例
        not_contain_sample_ratio: 对于不包含前3个token的歌曲，采样比例（更高概率）
        rng: 随机数生成器，如果为None则使用全局的np.random
        
    Returns:
        tuple: (是否应该采样, 是否在前3个) - 在前3个指position和duration都在前3个
    """
    # 获取position和duration列的唯一token
    position_idx = COLUMN_NAMES.index("position")
    duration_idx = COLUMN_NAMES.index("duration")
    
    position_tokens = get_data_tokens(data, position_idx)
    duration_tokens = get_data_tokens(data, duration_idx)
    
    # 判断是否在前3个
    position_in_top3 = bool(position_tokens & top_position_tokens)
    duration_in_top3 = bool(duration_tokens & top_duration_tokens)
    in_top3 = position_in_top3 and duration_in_top3
    
    if rng is None:
        rng = np.random
    
    # 条件1: 如果position或duration不包含前3个token，以更高概率采样
    if not position_in_top3 or not duration_in_top3:
        should_sample = rng.random() < not_contain_sample_ratio
        return should_sample, False
    
    # 条件2: 如果包含前3个token，则以固定百分比采样
    should_sample = rng.random() < contain_sample_ratio
    return should_sample, True


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 resample_songs(file_paths, top_position_tokens, top_duration_tokens,
                   contain_sample_ratio=0.3, not_contain_sample_ratio=0.9, 
                   num_threads=None, seed=42):
    """
    根据新逻辑进行重采样
    
    Args:
        file_paths: 文件路径列表
        top_position_tokens: position列的前3个最常见token集合
        top_duration_tokens: duration列的前3个最常见token集合
        contain_sample_ratio: 对于包含前3个token的歌曲，采样比例
        not_contain_sample_ratio: 对于不包含前3个token的歌曲，采样比例（更高概率）
        num_threads: 线程数，None表示使用CPU核心数
        seed: 随机种子
        
    Returns:
        tuple: (sampled_paths, sampled_indices, stats)
    """
    import threading
    
    # 为每个线程创建独立的随机数生成器
    thread_local = threading.local()
    
    def get_thread_rng():
        """获取当前线程的随机数生成器"""
        if not hasattr(thread_local, 'rng'):
            # 使用线程ID和种子创建独立的随机数生成器
            thread_id = threading.current_thread().ident
            thread_local.rng = np.random.RandomState(seed + hash(thread_id) % 1000000)
        return thread_local.rng
    
    if num_threads is None:
        num_threads = min(mp.cpu_count(), len(file_paths))
    
    print(f"\n开始重采样，使用 {num_threads} 个线程...")
    print(f"  包含前3个token的采样比例: {contain_sample_ratio*100:.1f}%")
    print(f"  不包含前3个token的采样比例: {not_contain_sample_ratio*100:.1f}%")
    
    def _should_sample_wrapper(args_tuple):
        """判断是否采样的包装函数（用于多线程）"""
        file_path, top_pos, top_dur, contain_ratio, not_contain_ratio = args_tuple
        try:
            # 使用线程本地的随机数生成器
            thread_rng = get_thread_rng()
            should_sample, in_top3 = should_sample_song(
                file_path, top_pos, top_dur, contain_ratio, not_contain_ratio, thread_rng
            )
            return (file_path, should_sample, in_top3, None)
        except Exception as e:
            return (file_path, False, False, str(e))
    
    # 准备参数
    sample_args = [(file_path, top_position_tokens, top_duration_tokens, 
                    contain_sample_ratio, not_contain_sample_ratio)
                   for file_path in file_paths]
    
    # 使用多线程判断每首歌是否应该采样
    sampled_paths = []
    sampled_indices = []
    stats = {
        'not_in_top3_count': 0,      # 不在前3个的歌曲数量
        'not_in_top3_sampled': 0,     # 不在前3个且被采样的歌曲数量
        'in_top3_count': 0,           # 在前3个的歌曲数量
        'in_top3_sampled': 0          # 在前3个且被采样的歌曲数量
    }
    
    # 限制并发任务数量，避免一次性提交过多任务
    batch_size = min(1000, len(file_paths))
    results = {}
    
    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        # 分批提交任务
        for batch_start in range(0, len(sample_args), batch_size):
            batch_end = min(batch_start + batch_size, len(sample_args))
            batch_args = sample_args[batch_start:batch_end]
            
            futures = {executor.submit(_should_sample_wrapper, args): args[0]
                       for args in batch_args}
            
            # 收集结果
            for future in tqdm(as_completed(futures), total=len(futures), 
                             desc=f"判断采样 [{batch_start+1}-{batch_end}/{len(file_paths)}]",
                             leave=False):
                try:
                    file_path, should_sample, in_top3, error = future.result(timeout=60)
                    results[file_path] = (should_sample, in_top3, error)
                    if error:
                        print(f"警告: 处理 {file_path} 时出错: {error}")
                except Exception as e:
                    print(f"错误: 获取结果时出错: {e}")
    
    # 按原始顺序处理结果，并统计
    for idx, file_path in enumerate(file_paths):
        if file_path not in results:
            continue
        
        should_sample, in_top3, error = results[file_path]
        if error:
            continue
        
        # 统计信息
        if in_top3:
            stats['in_top3_count'] += 1
            if should_sample:
                stats['in_top3_sampled'] += 1
        else:
            stats['not_in_top3_count'] += 1
            if should_sample:
                stats['not_in_top3_sampled'] += 1
        
        if should_sample:
            sampled_paths.append(file_path)
            sampled_indices.append(idx)
    
    print(f"\n采样完成:")
    print(f"  总歌曲数: {len(file_paths)}")
    print(f"  采样歌曲数: {len(sampled_paths)}")
    print(f"  采样比例: {len(sampled_paths)/len(file_paths)*100:.2f}%")
    print(f"  不在前3个的歌曲数: {stats['not_in_top3_count']}")
    print(f"  不在前3个且被采样的歌曲数: {stats['not_in_top3_sampled']}")
    if stats['not_in_top3_count'] > 0:
        print(f"  不在前3个的歌曲采样比例: {stats['not_in_top3_sampled']/stats['not_in_top3_count']*100:.2f}%")
    print(f"  在前3个的歌曲数: {stats['in_top3_count']}")
    print(f"  在前3个且被采样的歌曲数: {stats['in_top3_sampled']}")
    if stats['in_top3_count'] > 0:
        print(f"  在前3个的歌曲采样比例: {stats['in_top3_sampled']/stats['in_top3_count']*100:.2f}%")
    
    return sampled_paths, sampled_indices, stats


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="基于position和duration token的重采样V2")
    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_v2",
                       help="输出目录路径")
    parser.add_argument("--contain_sample_ratio", type=float, default=0.1,
                       help="对于包含前3个token的歌曲，采样比例 (默认: 0.1)")
    parser.add_argument("--not_contain_sample_ratio", type=float, default=0.9,
                       help="对于不包含前3个token的歌曲，采样比例 (默认: 0.9)")
    parser.add_argument("--top_k", type=int, default=3,
                       help="使用前k个最常见的token (默认: 3)")
    parser.add_argument("--seed", type=int, default=42,
                       help="随机种子")
    parser.add_argument("--num_threads", type=int, default=None,
                       help="线程数，None表示使用CPU核心数 (默认: None)")
    
    args = parser.parse_args()
    
    # 1. 加载position和duration的前top_k个最常见token
    top_position_tokens = load_top_tokens_from_json(
        args.json_path, "position", top_k=args.top_k
    )
    top_duration_tokens = load_top_tokens_from_json(
        args.json_path, "duration", top_k=args.top_k
    )
    
    if not top_position_tokens or not top_duration_tokens:
        print("错误: 无法加载前top_k个token")
        return
    
    # 2. 获取所有数据文件路径
    file_paths = get_data_file_paths(args.data_dir)
    
    if not file_paths:
        print("错误: 未找到任何数据文件")
        return
    
    print(f"\n共找到 {len(file_paths)} 个数据文件")
    
    # 3. 根据新逻辑进行重采样
    sampled_paths, sampled_indices, stats = resample_songs(
        file_paths,
        top_position_tokens,
        top_duration_tokens,
        contain_sample_ratio=args.contain_sample_ratio,
        not_contain_sample_ratio=args.not_contain_sample_ratio,
        num_threads=args.num_threads,
        seed=args.seed
    )
    
    # 4. 保存采样结果（延迟加载）
    sampled_data = sampled_paths  # 使用路径，延迟加载
    save_sampled_data(sampled_data, sampled_paths, args.output_dir, num_threads=args.num_threads)
    
    # 5. 保存采样索引（可选，用于后续分析）
    indices_file = Path(args.output_dir) / "sampled_indices.npy"
    np.save(indices_file, np.array(sampled_indices))
    print(f"\n采样索引已保存到: {indices_file}")
    
    # 6. 保存采样信息
    info = {
        "total_samples": len(file_paths),
        "sampled_samples": len(sampled_paths),
        "contain_sample_ratio": args.contain_sample_ratio,
        "not_contain_sample_ratio": args.not_contain_sample_ratio,
        "top_k": args.top_k,
        "top_position_tokens": sorted(list(top_position_tokens)),
        "top_duration_tokens": sorted(list(top_duration_tokens)),
        "stats": stats
    }
    
    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()

