import os
import zarr
import dask.array as da
from tqdm import tqdm
import yaml
#from numcodecs import Blosc #hua add to replace zarr.Blosc

def main():
    # Load config
    with open('config.yaml', 'r') as f:
        config = yaml.safe_load(f)

    input_dir = config['inference_config']['output']['path']
    output_path = config['merge_inference_config']['output']['path']

    celltype_list, cap_list = load_inputs(config)

    print(celltype_list)

    for celltype in celltype_list:
        merge_cap_predictions(input_dir, output_path, celltype, cap_list)

def merge_cap_predictions(input_dir, output_path, celltype, cap_list):
    input_dir = os.path.join(input_dir, celltype)
    output_path = os.path.join(output_path, celltype)
    os.makedirs(output_path, exist_ok=True)

    cap_list = sorted(cap_list)

    # Save cap list to file
    with open(os.path.join(output_path, f'cap_list.txt'), 'w') as f:
        for cap in cap_list:
            f.write(f"{cap}\n")

    # Construct paths to zarr files for each CAP
    zarr_files = []
    for cap in cap_list:
        zarr_path = os.path.join(input_dir, cap, 'processed', 'data.zarr')
        if os.path.exists(zarr_path):
            zarr_files.append(zarr_path)
        else:
            print(f"Warning: {zarr_path} does not exist, skipping")
    
    if not zarr_files:
        raise ValueError("No valid zarr files found to merge")

    # Get chrs from first zarr file
    infer_chrs = zarr.open(zarr_files[0], mode='r')['chrs']
    infer_chrs = [c for c in infer_chrs]

    # 3. Open or create output Zarr store (directory store on disk in this case)
    output_file = os.path.join(output_path, f'data.zarr')
    root = zarr.group(store = output_file, overwrite = False) # init zarr group

    # 4. Iterate over each chromosome to process
    
    for chrom in tqdm(infer_chrs, desc="Processing chromosomes"):
        files = [os.path.join(input_dir, cap, 'processed', 'data.zarr', f'chrs/{chrom}') for cap in cap_list]
        # Load each 1D Zarr file as a Dask array (lazy loaded, no data read yet)
        arrays = []
        for i, f in enumerate(tqdm(files, desc=f"Loading arrays for {chrom}", leave=False)):
            arrays.append(da.from_zarr(f))
        # (All arrays must have the same length for stacking; ensure shapes align)
        
        # 5. Stack arrays along a new axis (axis=0 creates shape (n_files, length))
        print(f"Stacking {len(arrays)} arrays for {chrom}...")
        # Stack arrays along the second axis (axis=1) instead of the first
        # This will create a shape of (length, n_files) instead of (n_files, length)
        stacked = da.stack(arrays, axis=1)
        
        # 6. Rechunk to have chunk size appropriate for the new shape (1_000_000, n)
        n_files = len(files)
        cap_chunk_size = min(n_files, 100)
        chunk_size = (1_000_000, cap_chunk_size)
        print(f"Rechunking array for {chrom} with chunk size {chunk_size}...")
        stacked = stacked.rechunk(chunks=chunk_size)
        
        # 7. Write this stacked array to the Zarr store under chrs/<chrom>
        # Use component to specify the subgroup path, and allow overwrite if exists.
        print(f"Writing {chrom} to zarr store...")
        da.to_zarr(stacked, output_file, component=f"chrs/{chrom}", overwrite=True, compressor=zarr.Blosc(cname='zstd', clevel=3, shuffle=zarr.Blosc.SHUFFLE))

        print(f"Completed {chrom}: shape {stacked.shape}, chunk size {stacked.chunksize}")

def load_inputs(config):
    celltype_list = read_list(os.path.join(config['inference_config']['input']['root'], config['inference_config']['input']['celltype_list_path']))
    cap_list = read_list(os.path.join(config['inference_config']['input']['root'], config['inference_config']['input']['cap_list_path']))
    return celltype_list, cap_list

def read_list(list_path):
    item_list = []
    with open(list_path, 'r') as file:
        for line in file:
            item_list.append(line.strip())
    return item_list
    
if __name__ == '__main__':
    main()