import os
import numpy
import cv2
import matplotlib.pyplot as plt

from model.model import Model
from rbm.entity import Entity, EntityParams, TrainingParams
from image.sub_image import SubImageExtract, normalize
from rbm.matrix import Mat, np, convert
from rbm.status import CheckpointStatus

DATA_DIR  = '/media/jens/cifs/bilder/MachineVision/Caltech_WebFaces/'
PATCH     = 8
STRIDE    = 4    # 50% overlap; set equal to PATCH for non-overlapping
GRAYSCALE = False  # reassigned in __main__ when --grayscale is set
N_CH      = 1 if GRAYSCALE else 3
N_VIS     = N_CH * PATCH * PATCH   # 1024 grayscale / 3072 colour
N_HID     = 64
N_IMAGES  = 50


class TestModel(Model):
	def __init__(self, name: str, work_dir: str = '.', l1_lambda: float = 0.0,
	             num_epochs: int = 1000, mini_batch_size: int = 1000):
		super().__init__(name, work_dir)
		self.unit1 = Entity(
			(N_VIS, N_HID),
			EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False),
			TrainingParams(learning_rate=0.001, momentum=0.9, num_epochs=num_epochs,
			               mini_batch_size=mini_batch_size, l1_lambda=l1_lambda)
		)

	def forward(self, x: Mat) -> Mat:
		return self.unit1.forward(x)

	def reconstruct(self, x: Mat) -> Mat:
		return self.unit1.reconstruct(x)


def pad_to_stride(img: numpy.ndarray) -> numpy.ndarray:
	"""Pad so (H - PATCH) and (W - PATCH) are divisible by STRIDE (required by SubImageExtract.check)."""
	h, w = img.shape[:2]
	ph = (-h) % STRIDE if h >= PATCH else PATCH - h
	pw = (-w) % STRIDE if w >= PATCH else PATCH - w
	return numpy.pad(img, ((0, ph), (0, pw), (0, 0)), mode='constant')


def extract_patches_numpy(img_chw: numpy.ndarray) -> numpy.ndarray:
	"""Extract overlapping patches (stride=STRIDE) from (C,H,W) array; returns (N, C*P*P)."""
	C, H, W = img_chw.shape
	rows = [img_chw[:, x:x+PATCH, y:y+PATCH].flatten()
	        for x in range(0, H - PATCH + 1, STRIDE)
	        for y in range(0, W - PATCH + 1, STRIDE)]
	return numpy.stack(rows, axis=0)


def _img_to_chw(img: numpy.ndarray) -> numpy.ndarray:
	"""Convert cv2 BGR image to (C, H, W) float64 array, respecting GRAYSCALE flag."""
	if GRAYSCALE:
		gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).astype(numpy.float64) / 255.0
		return gray[numpy.newaxis]                                    # (1, H, W)
	img_f = img[:, :, ::-1].astype(numpy.float64) / 255.0           # BGR→RGB
	return numpy.ascontiguousarray(numpy.transpose(img_f, (2, 0, 1)))  # (3, H, W)


def load_patches(data_dir: str, n_images: int = N_IMAGES, start: int = 0):
	"""Load images, extract overlapping patches, return normalised cupy array."""
	all_patches = []
	files = sorted(f for f in os.listdir(data_dir) if f.endswith('.jpg'))[start:start + n_images]
	for fname in files:
		img = cv2.imread(os.path.join(data_dir, fname))
		if img is None or img.shape[0] < PATCH or img.shape[1] < PATCH:
			continue
		img = pad_to_stride(img)
		img_chw = _img_to_chw(img)
		all_patches.append(extract_patches_numpy(img_chw))

	patches_np = numpy.concatenate(all_patches, axis=0)
	std = numpy.std(patches_np, axis=1)
	patches_np = patches_np[std > 0.01]

	# Normalize on CPU; return numpy — train() converts mini-batches to GPU on the fly
	mean = patches_np.mean(axis=1, keepdims=True)
	std2 = patches_np.std(axis=1, keepdims=True)
	return (patches_np - mean) / numpy.where(std2 < 1e-8, 1.0, std2)


def load_image_patches(path: str):
	"""Extract normalised patches from a single image via SubImageExtract; return (patches, nx_steps, ny_steps)."""
	sub = SubImageExtract(PATCH, PATCH, STRIDE, STRIDE)
	img = cv2.imread(path)
	img = pad_to_stride(img)
	h, w = img.shape[:2]
	nx_steps = (h - PATCH) // STRIDE + 1
	ny_steps = (w - PATCH) // STRIDE + 1

	img_nchw = np.array(_img_to_chw(img)[numpy.newaxis])
	patches_nchw = sub(img_nchw)
	patches = patches_nchw.reshape(-1, N_VIS)
	return normalize(patches), nx_steps, ny_steps


def assemble_overlapping(patch_list, nx_steps: int, ny_steps: int) -> numpy.ndarray:
	"""Reconstruct image from overlapping patches by averaging contributions per pixel."""
	H = (nx_steps - 1) * STRIDE + PATCH
	W = (ny_steps - 1) * STRIDE + PATCH
	accum = numpy.zeros((N_CH, H, W), dtype=numpy.float64)
	count = numpy.zeros((1,    H, W), dtype=numpy.float64)

	for k, p in enumerate(patch_list):
		x = (k // ny_steps) * STRIDE
		y = (k %  ny_steps) * STRIDE
		arr = convert(p).reshape(N_CH, PATCH, PATCH)
		accum[:, x:x+PATCH, y:y+PATCH] += arr
		count[0, x:x+PATCH, y:y+PATCH] += 1.0

	img = (accum / numpy.maximum(count, 1)).transpose(1, 2, 0).squeeze()  # (H,W,C) or (H,W)
	lo, hi = img.min(), img.max()
	return numpy.clip((img - lo) / (hi - lo + 1e-8), 0, 1)


def show_filters(entity: Entity, rows: int = 8, cols: int = 16):
	W = convert(entity.state.w_hv)
	fig, axes = plt.subplots(rows, cols, figsize=(cols * 1.2, rows * 1.2))
	fig.suptitle(f'Learned GB-RBM filters  ({N_HID} hidden units, {PATCH}×{PATCH} RGB patches)', fontsize=10)
	for j in range(rows * cols):
		ax = axes[j // cols][j % cols]
		ax.axis('off')
		if j >= N_HID:
			continue
		filt = W[:, j].reshape(N_CH, PATCH, PATCH).transpose(1, 2, 0).squeeze()
		lo, hi = filt.min(), filt.max()
		ax.imshow(numpy.clip((filt - lo) / (hi - lo + 1e-8), 0, 1),
		          cmap='gray' if GRAYSCALE else None)
	plt.tight_layout()


def show_reconstructions(model: TestModel, patches, n_show: int = 10):
	fig, axes = plt.subplots(2, n_show, figsize=(n_show * 1.5, 3.5))
	axes[0, 0].set_ylabel('Original')
	axes[1, 0].set_ylabel('Recon')
	fig.suptitle('Patch reconstructions (normalised display)', fontsize=10)
	err_total = 0.0

	# Convert to GPU if patches arrived as a CPU numpy array
	if isinstance(patches, numpy.ndarray):
		patches = np.array(patches)

	cmap = 'gray' if GRAYSCALE else None

	def to_img(p):
		a = convert(p).reshape(N_CH, PATCH, PATCH).transpose(1, 2, 0).squeeze()
		return numpy.clip((a - a.min()) / (a.max() - a.min() + 1e-8), 0, 1)

	for i in range(n_show):
		inp   = patches[i]
		recon = model.reconstruct(model.forward(inp))
		err_total += float(np.mean(np.abs(inp - recon)))
		axes[0, i].imshow(to_img(inp),   cmap=cmap); axes[0, i].axis('off')
		axes[1, i].imshow(to_img(recon), cmap=cmap); axes[1, i].axis('off')
	print(f'Mean patch reconstruction MAE: {err_total / n_show:.4f}')
	plt.tight_layout()


def show_image_reconstruction(model: TestModel, img_path: str):
	patches, nx_steps, ny_steps = load_image_patches(img_path)
	recons = [model.reconstruct(model.forward(patches[i])) for i in range(patches.shape[0])]

	orig_grid  = assemble_overlapping([patches[i] for i in range(len(recons))], nx_steps, ny_steps)
	recon_grid = assemble_overlapping(recons, nx_steps, ny_steps)

	fig, axes = plt.subplots(1, 2, figsize=(12, 6))
	cmap = 'gray' if GRAYSCALE else None
	axes[0].imshow(orig_grid,  cmap=cmap); axes[0].set_title('Original (normalised)');        axes[0].axis('off')
	axes[1].imshow(recon_grid, cmap=cmap); axes[1].set_title('Reconstructed (overlap blended)'); axes[1].axis('off')
	fig.suptitle(os.path.basename(img_path), fontsize=9)
	plt.tight_layout()


if __name__ == '__main__':
	from argparse import ArgumentParser
	ap = ArgumentParser()
	ap.add_argument('--load_model', type=lambda s: s.lower() != 'false', default=True,
	                help='Load saved model weights (default: true)')
	ap.add_argument('--do_train',   type=lambda s: s.lower() != 'false', default=False,
	                help='Train the model (default: false)')
	ap.add_argument('--grayscale',  action='store_true', default=False,
	                help='Use single-channel grayscale patches (default: false)')
	ap.add_argument('--l1_lambda',  type=float, default=0.0,
	                help='L1 regularisation strength (default: 0.0)')
	ap.add_argument('--num_epochs',      type=int,   default=1000,
	                help='Number of training epochs (default: 1000)')
	ap.add_argument('--mini_batch_size', type=int,   default=1000,
	                help='Mini-batch size (default: 1000)')
	ap.add_argument('--n_images',        type=int,   default=N_IMAGES,
	                help=f'Number of training images (default: {N_IMAGES})')
	args = ap.parse_args()

	if args.grayscale:
		GRAYSCALE = True
		N_CH      = 1
		N_VIS     = PATCH * PATCH   # 1024

	prj_name = 'faces_sub_image_gray' if args.grayscale else 'faces_sub_image'
	work_dir = 'results'

	model = TestModel(prj_name, work_dir, l1_lambda=args.l1_lambda,
	                  num_epochs=args.num_epochs, mini_batch_size=args.mini_batch_size)
	model.init(0.001)

	if args.load_model:
		model.load()

	if args.do_train:
		print(f'Loading {args.n_images} training images (stride={STRIDE})...')
		train_patches = load_patches(DATA_DIR, args.n_images)
		print(f'Training on {train_patches.shape[0]} patches ({N_VIS}-dim each)')
		model.train(train_patches, status=CheckpointStatus(save_fn=model.save))
		model.save()

	# Show learned weight filters
	show_filters(model.unit1)

	# Show patch-level reconstructions on held-out images
	print('Loading test patches...')
	test_patches = load_patches(DATA_DIR, n_images=10, start=N_IMAGES)
	if test_patches.shape[0] < 10:
		test_patches = load_patches(DATA_DIR, n_images=10)
	show_reconstructions(model, test_patches)

	# Show full image reconstruction with overlap blending
	test_img = os.path.join(DATA_DIR, sorted(os.listdir(DATA_DIR))[60])
	show_image_reconstruction(model, test_img)

	plt.show()