import torch
from torch.utils.data import DataLoader
import numpy as np
import imageio
import os
import shutil
from helpers.utils import is_main_process, get_rank, get_world_size, safe_barrier


def delete_content_of_dir(folder):
    os.makedirs(folder, exist_ok=True)
    for filename in os.listdir(folder):
        file_path = os.path.join(folder, filename)
        try:
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as e:
            print('Failed to delete %s. Reason: %s' % (file_path, e))

def get_sample_for_visualization(data, preprocess_fn, num, dataset):
    for x in DataLoader(data, batch_size=num):
        break
    orig_image = (x[0] * 255.0).to(torch.uint8).permute(0, 2, 3, 1) if dataset == 'ffhq_1024' else x[0]
    preprocessed = preprocess_fn(x)[0]
    return orig_image, preprocessed


def _resolve_snoise(sampler, batch_size, use_snoise):
    if use_snoise and hasattr(sampler, 'snoise_tmp'):
        return [s[:batch_size].normal_() for s in sampler.snoise_tmp]
    if hasattr(sampler, 'neutral_snoise'):
        return [s[:batch_size] for s in sampler.neutral_snoise]
    return None



def generate_for_NN(sampler, orig, initial, *args):
    if len(args) == 4:
        snoise = None
        shape, ema_imle, fname, logprint = args
    elif len(args) == 5:
        snoise, shape, ema_imle, fname, logprint = args
    else:
        raise TypeError("generate_for_NN expected 4 or 5 trailing args")

    mb = shape[0]
    initial = initial[:mb].cuda()
    nns = sampler.sample(initial, ema_imle, snoise)
    batches = [orig[:mb], nns]
    n_rows = len(batches)
    im = np.concatenate(batches, axis=0).reshape((n_rows, mb, *shape[1:])).transpose([0, 2, 1, 3, 4]).reshape(
        [n_rows * shape[1], mb * shape[2], 3])

    logprint(f'printing samples to {fname}')
    imageio.imwrite(fname, im)


def generate_images_initial(H, sampler, orig, initial, snoise, shape, imle, ema_imle, fname, logprint, experiment=None):
    mb = shape[0]
    initial = initial[:mb]
    batches = [orig[:mb], sampler.sample(initial, imle, snoise)]

    temp_latent_rnds = torch.randn([mb, H.latent_dim], dtype=torch.float32).cuda()
    for t in range(H.num_rows_visualize + 4):
        temp_latent_rnds.normal_()
        tmp_snoise = _resolve_snoise(sampler, mb, H.use_snoise)
        batches.append(sampler.sample(temp_latent_rnds, imle, tmp_snoise))

    # if(H.use_snoise == True):
    #     tmp_snoise = [s[:mb].normal_() for s in sampler.snoise_tmp]
    # else:
    #     tmp_snoise = [s[:mb] for s in sampler.neutral_snoise]
    # batches.append(sampler.sample(temp_latent_rnds, imle, tmp_snoise))

    # if(H.use_snoise == True):
    #     tmp_snoise = [s[:mb].normal_() for s in sampler.snoise_tmp]
    # else:
    #     tmp_snoise = [s[:mb] for s in sampler.neutral_snoise]
    # batches.append(sampler.sample(temp_latent_rnds, imle, tmp_snoise))

    # tmp_snoise = [s[:mb] for s in sampler.neutral_snoise]
    # batches.append(sampler.sample(temp_latent_rnds, imle, tmp_snoise))

    # tmp_snoise = [s[:mb] for s in sampler.neutral_snoise]
    # temp_latent_rnds.normal_()
    # batches.append(sampler.sample(temp_latent_rnds, imle, tmp_snoise))

    n_rows = len(batches)
    im = np.concatenate(batches, axis=0).reshape((n_rows, mb, *shape[1:])).transpose([0, 2, 1, 3, 4]).reshape(
        [n_rows * shape[1], mb * shape[2], 3])

    logprint(f'printing samples to {fname}')
    imageio.imwrite(fname, im)
    if(experiment):
        experiment.log_image(fname, overwrite=True)


def generate_visualization(H, sampler, orig, initial, last_initial, latent_for_visualization,
                           shape, imle, fname, logprint, experiment=None):
    _ = last_initial
    _ = latent_for_visualization
    generate_images_initial(H, sampler, orig, initial, None, shape, imle, imle, fname, logprint, experiment)

def generate_and_save(H, imle, sampler, n_samp, subdir='fid'):
    rank = get_rank()
    world_size = get_world_size()

    if is_main_process():
        delete_content_of_dir(f'{H.save_dir}/{subdir}')
    safe_barrier()

    per_rank = (n_samp + world_size - 1) // world_size
    start_idx = rank * per_rank
    end_idx = min(start_idx + per_rank, n_samp)
    local_n = end_idx - start_idx

    # Inference-time latent-noise amplitude (1.0 = standard N(0, I)).
    latent_std = float(getattr(H, 'eval_latent_std', 1.0) or 1.0)

    with torch.no_grad():
        temp_latent_rnds = torch.randn([H.imle_batch, H.latent_dim], dtype=torch.float32).cuda()
        generated = 0
        while generated < local_n:
            batch_size = min(H.imle_batch, local_n - generated)

            temp_latent_rnds.normal_(mean=0.0, std=latent_std)
            tmp_snoise = _resolve_snoise(sampler, H.imle_batch, H.use_snoise)
            samp = sampler.sample(temp_latent_rnds, imle, tmp_snoise)

            for j in range(batch_size):
                imageio.imwrite(f'{H.save_dir}/{subdir}/{start_idx + generated + j}.png', samp[j])
            generated += batch_size
    safe_barrier()