import json, sys, os, numpy as np, h5py, re
from pathlib import Path
from datetime import datetime, timedelta
from tqdm import tqdm

INDEX_PATH = Path("/mnt/kingston/seismic_webapp/data/index.json")
OUTPUT_DIR = Path("/mnt/kingston/seismic_webapp/data/rms_cache")
SAMPLE_RATE = 200

def fix_path(p):
    p = p.replace('\\', '/')
    if p.startswith('F:/'): return '/mnt/kingston/' + p[3:]
    if p.startswith('E:/'): return '/mnt/data_sdb1/' + p[3:]
    return p

def compute_rms(h5_path):
    h5_path = fix_path(h5_path)
    if not os.path.exists(h5_path): return None
    try:
        # Extraire le timestamp réel du nom de fichier (Julian day)
        match = re.search(r'auto_(\d+)_(\d{6})_b', os.path.basename(h5_path))
        if not match: return None
        julian, time_str = int(match.group(1)), match.group(2)
        h, m, s = int(time_str[:2]), int(time_str[2:4]), int(time_str[4:6])
        start_ts = int(datetime(2020, 1, 1).timestamp() + (julian - 1) * 86400 + h * 3600 + m * 60 + s)
        
        with h5py.File(h5_path, 'r') as f:
            ds = f['adc_values']
            samples = ds[0:5000]
            rms = float(np.sqrt(np.mean(samples.astype(np.float64)**2)))
            return [{'ts': start_ts, 'rms': rms}]
    except: return None

def main():
    with open(INDEX_PATH, 'r') as f: index = json.load(f)
    nodes = index.get('nodes', {})
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    
    for date in index['dates']:
        channel = "ch0"
        output_file = OUTPUT_DIR / f"rms_{{date}}_{{channel}}.json"
        print(f"Processing {date}...")
        results = {}
        for nid, node in tqdm(nodes.items(), desc=f"Nodes {date}"):
            files = node.get('files', [])
            # Filtrer les fichiers par Julian Day correspondant à la date
            dt = datetime.strptime(date, '%Y-%m-%d')
            target_julian = dt.timetuple().tm_yday
            
            target = next((f for f in files if f['julian'] == target_julian and f'_{{channel}}_' in f['path']), None)
            if not target and files: 
                target = next((f for f in files if f['julian'] == target_julian), None)
            
            if target:
                data = compute_rms(target['path'])
                if data: results[nid] = data
        
        if results:
            with open(output_file, 'w') as f:
                json.dump({'date':date, 'channel':channel, 'nodes':results}, f)

if __name__ == '__main__': main()