""" Main script to generate samples of FINAL data 

	Just a stack of {w2, x, y} samples from the final model 

"""
import torch
import pickle
from torch.utils.data import DataLoader
import argparse
import os
from tqdm.auto import tqdm

from cfg.dataloader_pickle import PickleDataset
from cfg.embedding import JointEmbedding2
from cfg.unet import Unet
from cfg.utils import get_named_beta_schedule
from cfg.diffusion import GaussianDiffusion
from path_constant import project_root


# =============================================
# =           Loading blocks                  =
# =============================================



def load_data(pkl_loc, batch_size):
	""" Loads the observational data from a pickle file """
	dataset = PickleDataset(pkl_file=pkl_loc)
	loader = DataLoader(dataset, num_workers=8, batch_size=batch_size,
						shuffle=False, drop_last=False)
	return loader


def load_model(diffuser_loc, device, w=0.0):
	""" Loads the architecture for the diffusion model we've already trained """
	net = Unet(
		in_ch=3,
		mod_ch=64,
		out_ch=3,
		ch_mul=[1,2,2,2],
		num_res_blocks=2,
		cdim=64, 
		use_conv=True,
		droprate=0,
		dtype=torch.float32).to(device)

	checkpoint = torch.load(diffuser_loc, map_location='cpu')
	out = net.load_state_dict(checkpoint['net'])
	assert len(out.missing_keys) + len(out.unexpected_keys) == 0

	betas = get_named_beta_schedule(num_diffusion_timesteps=1000)
	diffusion = GaussianDiffusion(
					dtype=torch.float32,
					model=net,
					betas=betas,
					w=w, # truly conditional sampling with no upweighting
					v=1.0,
					device=device)
	cemblayer = JointEmbedding2(num_labels_0=10, num_labels_1=2,
                               d_model=64, channels=3,
                               dim=64, hw=32).to(device)
	out = cemblayer.load_state_dict(checkpoint['cemblayer'])
	assert len(out.missing_keys) + len(out.unexpected_keys) == 0

	diffusion.model.eval()
	cemblayer.eval()
	return {'diffusion': diffusion, 'cemblayer': cemblayer}



# ========================================================
# =           Sample 1/multiple batches                  =
# ========================================================
@torch.no_grad()
def sample_from_batch(batch, diffuser, device, ddim=True, drop_label=True):

	bsz = batch['X'].shape[0]

	# 1: Get uniformly random w2
	#w2a = batch['W2a'].to(device)
	w2a = torch.randint(low=0, high=10, size=(bsz,)).to(device)
	#w2b = batch['W2b'].to(device)
	w2b = torch.randint(low=0, high=2, size=(bsz,)).to(device)

	# 2: Take x from data 
	x = batch['X'].to(device)

	# 3. Sample from trained P(y | x,w2)
	cemb = diffuser['cemblayer'](x, w2a, w2b, threshold=0.0)
	cemb_uncond = diffuser['cemblayer'](x, w2a, w2b, drop_label=bool(drop_label), drop_image=True, threshold=999.0)


	# lst= []
	# for iter in tqdm(range(20)):

	if ddim:
		generated = diffuser['diffusion'].ddim_sample((bsz, 3, 32, 32), 50, 0, 'linear',
													  cemb=cemb, cemb_uncond=cemb_uncond,disable_tqdm=True)
	else:
		generated = diffuser['diffusion'].sample((bsz, 3, 32, 32), cemb=cemb,
												 cemb_uncond=cemb_uncond, disable_tqdm=True)

		# lst.append(generated.unsqueeze(0))
	# y= torch.cat(lst, dim=0)
	y= generated

	# Return (w1, w2a, w2b, x, y)
	return (w2a.cpu(), w2b.cpu(), x.cpu(), y.cpu())





def sample_batches(dataloader, diffuser, n_samples, device, ddim=True, drop_label=True):

	data = {'W2a': [], 'W2b': [], 'X': [], 'Y': []}
	count = 0
	iterator = tqdm(total=n_samples)
	while count < n_samples:
		for batch in dataloader:
			w2a, w2b, x,y = sample_from_batch(batch, diffuser, device, ddim=ddim, drop_label=drop_label)
			data['W2a'].append(w2a)
			data['W2b'].append(w2b)
			data['X'].append(x)
			data['Y'].append(y)
			count += x.shape[0]
			print('Count-->',count)
			iterator.update(n=x.shape[0])
			if count >= n_samples:
				break

	return {k: torch.cat(v) for k,v in data.items()}


def save_datadict(data_dict, save_dir):
	os.makedirs(save_dir, exist_ok=True)
	save_loc = os.path.join(save_dir, 'final_W2XY.pkl')
	with open(save_loc, 'wb') as f:	
		pickle.dump(data_dict, f)


# ===========================================
# =           Main block                    =
# ===========================================


def main():
	parser = argparse.ArgumentParser(description='script to generate samples for retraining purposes')
	parser.add_argument('--pkl_loc', type=str, default=f"{project_root}/Baselines/Compare/baseline_samples/do_X.pkl", help='location of pickle training data file') # Base data here
	# parser.add_argument('--diffuser_loc', type=str, default="./final_model_NODROP/wEpoch1000/ckpt_1000_checkpoint.pt", help='location of saved diffusion model')
	parser.add_argument('--diffuser_loc', type=str, default="./final_model_NODROP/wEpoch300/ckpt_300_checkpoint.pt", help='location of saved diffusion model')
	parser.add_argument('--n_samples', type=int, default=6000, help='how many samples to generate')
	parser.add_argument('--batch_size', type=int, default=200, help='batch size')
	parser.add_argument('--device', type=int, default=1, help='which gpu to use')
	parser.add_argument('--ddim', type=int, default=1, help='1 if we want to use ddim sampling, 0 ow')
	parser.add_argument('--w', type=float, default=0.0)
	parser.add_argument('--drop_label', type=int, default=0)
	parser.add_argument('--save_dir', type=str, default=f"{project_root}/Baselines/Compare/baseline_samples/IDGEN", help='location of where to save results')

	params = parser.parse_args()
	params.ddim = bool(params.ddim)
	params.device = 'cuda:%s' % params.device


	dataloader = load_data(params.pkl_loc, params.batch_size)
	diffuser = load_model(params.diffuser_loc, params.device, w=params.w)
	data_dict = sample_batches(dataloader, diffuser, params.n_samples, params.device, params.ddim)
	save_datadict(data_dict, params.save_dir)

# Iterates through the input intervetion dataloader (here batch['X']) until n_samples are achieved.
# 	python3 gen_final_data.py --pkl_loc=base_data/napkin_mnist_train.pkl --diffuser_loc=final_model_NODROP/wEpoch1000/ckpt_1000_checkpoint.pt --n_samples=40 --device=1 --save_dir=final_data_Y_doX/
#  python3 gen_final_data.py --pkl_loc=base_data/napkin_mnist_train.pkl --diffuser_loc=final_model_NODROP/wEpoch1000/ckpt_1000_checkpoint.pt --n_samples=40 --batch_size=40 --device=1 --save_dir=final_data_Y_doX/
#  python3 gen_final_data.py --pkl_loc=baseline_samples/do_X.pkl --diffuser_loc=final_model_NODROP/wEpoch300/ckpt_300_checkpoint.pt --n_samples=20 --batch_size=200 --device=1 --save_dir=baseline_samples/
if __name__ == '__main__':
	main()