from spandrel import ModelLoader
import torch
from pathlib import Path
import gradio as App
import logging
import spaces
import time
import cv2
import os

from gradio import themes
from rich.console import Console
from rich.logging import RichHandler

from Scripts.SAD import GetDifferenceRectangles

# ============================== #
#          Core Settings         #
# ============================== #

Theme = themes.Citrus(
	primary_hue='blue',
	secondary_hue='blue',
	radius_size=themes.sizes.radius_xxl
).set(
	link_text_color='blue'
)
ModelDir = Path('./Models')
TempDir = Path('./Temp')
os.environ['GRADIO_TEMP_DIR'] = str(TempDir)
ModelFileType = '.pth'

# ============================== #
#            Logging             #
# ============================== #

logging.basicConfig(
	level=logging.INFO,
	format='%(message)s',
	datefmt='[%X]',
	handlers=[RichHandler(
		console=Console(),
		rich_tracebacks=True,
		omit_repeated_times=False,
		markup=True,
		show_path=False,
	)],
)
Logger = logging.getLogger('Video2x')
logging.getLogger('httpx').setLevel(logging.WARNING)

# ============================== #
#      Device Configuration      #
# ============================== #

@spaces.GPU
def GetDeviceName():
	Device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
	Logger.info(f'🧪 Using device: {str(Device).upper()}')
	return Device

Device = GetDeviceName()

# ============================== #
#       Utility Functions        #
# ============================== #

def HumanizeSeconds(Seconds):
	Hours = int(Seconds // 3600)
	Minutes = int((Seconds % 3600) // 60)
	Seconds = int(Seconds % 60)

	if Hours > 0:
		return f'{Hours}h {Minutes}m {Seconds}s'
	elif Minutes > 0:
		return f'{Minutes}m {Seconds}s'
	else:
		return f'{Seconds}s'

def HumanizedBytes(Size):
	Units = ['B', 'KB', 'MB', 'GB', 'TB']
	Index = 0
	while Size >= 1024 and Index < len(Units) - 1:
		Size /= 1024.0
		Index += 1
	return f'{Size:.2f} {Units[Index]}'

# ============================== #
#     Main Processing Logic      #
# ============================== #

@spaces.GPU
class Upscaler:
	def __init__(self):
		pass

	def ListModels(self):
		Models = sorted(
			[File.name for File in ModelDir.glob('*' + ModelFileType) if File.is_file()]
		)
		Logger.info(f'📚 Found {len(Models)} Models In Directory')
		return Models

	def LoadModel(self, ModelName):
		torch.cuda.empty_cache()
		Model = (
			ModelLoader()
			.load_from_file(ModelDir / (ModelName + ModelFileType))
			.to(Device)
			.eval()
		)
		Logger.info(f'🤖 Loaded Model {ModelName} Onto {str(Device).upper()}')
		return Model

	def UnloadModel(self):
		if Device.type == 'cuda':
			torch.cuda.empty_cache()
		Logger.info('🤖 Model Unloaded Successfully')

	def CleanUp(self):
		self.UnloadModel()
		Logger.info('🧹 Temporary Files Cleaned Up')

	def Process(self, InputVideo, InputModel, InputUseRegions, InputThreshold, InputMinPercentage, InputMaxRectangles, InputPadding, Progress=App.Progress()):
		if not InputVideo:
			Logger.warning('❌ No Video Provided')
			App.Warning('❌ No Video Provided')
			return None, None

		Progress(0, desc='⚙️ Loading Model')
		Model = self.LoadModel(InputModel)

		Logger.info(f'📼 Processing Video: {Path(InputVideo).name}')
		Progress(0, desc='📼 Processing Video')
		Video = cv2.VideoCapture(InputVideo)

		FrameRate = Video.get(cv2.CAP_PROP_FPS)
		FrameCount = int(Video.get(cv2.CAP_PROP_FRAME_COUNT))
		Width = int(Video.get(cv2.CAP_PROP_FRAME_WIDTH))
		Height = int(Video.get(cv2.CAP_PROP_FRAME_HEIGHT))

		Logger.info(f'📏 Video Properties: {FrameCount} Frames, {FrameRate} FPS, {Width}x{Height}')

		PerFrameProgress = 1 / FrameCount
		FrameProgress = 0.0
		StartTime = time.time()
		Times = []

		while True:
			Ret, Frame = Video.read()
			if not Ret:
				break

			FrameRgb = cv2.cvtColor(Frame, cv2.COLOR_BGR2RGB)
			FrameForTorch = FrameRgb.transpose(2, 0, 1)
			FrameForTorch = torch.from_numpy(FrameForTorch).unsqueeze(0).to(Device).float() / 255.0

			RetNext, NextFrame = Video.read()
			if not RetNext:
				NextFrame = Frame

			DiffResult = GetDifferenceRectangles(
				Frame,
				NextFrame,
				Threshold=InputThreshold,
				Rows=12,
				Columns=20,
				Padding=InputPadding

			)
			SimilarityPercentage = DiffResult['SimilarPercentage']
			Rectangles = DiffResult['Rectangles']

			if SimilarityPercentage > InputMinPercentage and len(Rectangles) < InputMaxRectangles and InputUseRegions:
				Logger.info(f'🟩 Frame {int(Video.get(cv2.CAP_PROP_POS_FRAMES))}: {SimilarityPercentage:.2f}% Similar, {len(Rectangles)} Regions To Upscale')
				Cols = DiffResult['Columns']
				Rows = DiffResult['Rows']
				FrameHeight, FrameWidth = Frame.shape[:2]
				SegmentWidth = FrameWidth // Cols
				SegmentHeight = FrameHeight // Rows
				for X, Y, W, H in Rectangles:
					X1 = X * SegmentWidth
					Y1 = Y * SegmentHeight
					X2 = FrameWidth if X + W == Cols else X1 + W * SegmentWidth
					Y2 = FrameHeight if Y + H == Rows else Y1 + H * SegmentHeight

					Region = Frame[Y1:Y2, X1:X2]
					RegionRgb = cv2.cvtColor(Region, cv2.COLOR_BGR2RGB)
					RegionTorch = torch.from_numpy(RegionRgb.transpose(2, 0, 1)).unsqueeze(0).to(Device).float() / 255.0
					UpscaledRegion = Model(RegionTorch)[0].cpu().numpy().transpose(1, 2, 0) * 255.0 # type: ignore
					UpscaledRegion = cv2.cvtColor(UpscaledRegion.astype('uint8'), cv2.COLOR_RGB2BGR)
					RegionHeight, RegionWidth = Region.shape[:2]
					UpscaledRegion = cv2.resize(UpscaledRegion, (RegionWidth, RegionHeight), interpolation=cv2.INTER_CUBIC)
					Frame[Y1:Y2, X1:X2] = UpscaledRegion
				OutputFrame = Frame
			else:
				Logger.info(f'🟥 Frame {int(Video.get(cv2.CAP_PROP_POS_FRAMES))}: {SimilarityPercentage:.2f}% Similar, Upscaling Full Frame')
				OutputFrame = Model(FrameForTorch)[0].cpu().numpy().transpose(1, 2, 0) * 255.0 # type: ignore
				OutputFrame = cv2.cvtColor(OutputFrame.astype('uint8'), cv2.COLOR_RGB2BGR)
				OutputFrame = cv2.resize(OutputFrame, (Width, Height), interpolation=cv2.INTER_CUBIC)

			CurrentFrameNumber = int(Video.get(cv2.CAP_PROP_POS_FRAMES))
			if Times:
				AverageTime = sum(Times) / len(Times)
				Eta = HumanizeSeconds((FrameCount - CurrentFrameNumber) * AverageTime)
			else:
				Eta = None

			Progress(FrameProgress, desc=f'📦 Processed Frame {len(Times)+1}/{FrameCount} - {Eta}')
			Logger.info(f'📦 Processed Frame {len(Times)+1}/{FrameCount} - {Eta}')

			cv2.imwrite(f'{TempDir}/Upscaled_Frame_{CurrentFrameNumber:05d}.png', OutputFrame)

			DeltaTime = time.time() - StartTime
			Times.append(DeltaTime)
			StartTime = time.time()
			FrameProgress += PerFrameProgress

		Progress(1, desc='📦 Cleaning Up')
		self.CleanUp()
		return InputVideo, InputVideo

# ============================== #
#          Streamlined UI        #
# ============================== #

with App.Blocks(
	title='Video Upscaler', theme=Theme, delete_cache=(-1, 1800)
) as Interface:
	App.Markdown('# 🎞️ Video Upscaler')
	App.Markdown('''
			  Space created by [Hyphonical](https://huggingface.co/Hyphonical), this space uses several models from [styler00dollar/VSGAN-tensorrt-docker](https://github.com/styler00dollar/VSGAN-tensorrt-docker/releases/tag/models)
			  You may always request adding more models by opening a [new discussion](https://huggingface.co/spaces/Hyphonical/Video2x/discussions/new). The main program uses spandrel to load the models and ffmpeg to process the video.
			  You may run out of time using the ZeroGPU, you could clone the space or run it locally for better performance.
			  ''')

	with App.Row():
		with App.Column():
			with App.Group():
				with App.Accordion(label='📜 Instructions', open=False):
					App.Markdown('''
						### How To Use The Video Upscaler
						1.  **Upload A Video:** Begin by uploading your video file using the 'Input Video' section.
						2.  **Select A Model:** Choose an appropriate upscaling model from the 'Select Model' dropdown menu.
						3.  **Adjust Settings (Optional):**
						Modify the 'Frame Rate' slider if you want to change the output video's frame rate.
						Adjust the 'Tile Grid Size' for memory optimization. Larger models might require a higher grid size, but processing could be slower.
						4.  **Start Processing:** Click the '🚀 Upscale Video' button to begin the upscaling process.
						5.  **Download The Result:** Once the process is complete, download the upscaled video using the '💾 Download Video' button.
						> Tip: If you get a CUDA out of memory error, try increasing the Tile Grid Size. This will split the image into smaller tiles for processing, which can help reduce memory usage.
					''')
				InputVideo = App.Video(
					label='Input Video', sources=['upload'], height=300
				)
				ModelList = Upscaler().ListModels()
				ModelNames = [Path(Model).stem for Model in ModelList]
				InputModel = App.Dropdown(
					choices=ModelNames,
					label='Select Model',
					value=ModelNames[0],
				)
				with App.Accordion(label='⚙️ Advanced Settings', open=False):
					with App.Group():
						InputUseRegions = App.Checkbox(
							label='Use Regions',
							value=False,
							info='Use regions to upscale only the different parts of the video (⚡️ Experimental, Faster)',
							interactive=True
						)
						InputThreshold = App.Slider(
							label='Threshold',
							value=5,
							minimum=0,
							maximum=20,
							step=0.5,
							info='Threshold for the SAD algorithm to detect different regions',
							interactive=False
						)
						InputPadding = App.Slider(
							label='Padding',
							value=1,
							minimum=0,
							maximum=5,
							step=1,
							info='Extra padding to include neighboring pixels in the SAD algorithm',
							interactive=False
						)
						InputMinPercentage = App.Slider(
							label='Min Percentage',
							value=70,
							minimum=0,
							maximum=100,
							step=1,
							info='Minimum percentage of similarity to consider upscaling the full frame',
							interactive=False
						)
						InputMaxRectangles = App.Slider(
							label='Max Rectangles',
							value=8,
							minimum=1,
							maximum=10,
							step=1,
							info='Maximum number of rectangles to consider upscaling the full frame',
							interactive=False
						)
			SubmitButton = App.Button('🚀 Upscale Video')

		with App.Column(show_progress=True):
			with App.Group():
				OutputVideo = App.Video(
					label='Output Video', height=300, interactive=False, format=None
				)
			OutputDownload = App.DownloadButton(
				label='💾 Download Video', interactive=False
			)

	def ToggleRegionInputs(UseRegions):
		return (
			App.update(interactive=UseRegions),
			App.update(interactive=UseRegions),
			App.update(interactive=UseRegions),
			App.update(interactive=UseRegions),
		)

	InputUseRegions.change(
		fn=ToggleRegionInputs,
		inputs=[InputUseRegions],
		outputs=[InputThreshold, InputMinPercentage, InputMaxRectangles, InputPadding],
	)

	SubmitButton.click(
		fn=Upscaler().Process,
		inputs=[
			InputVideo,
			InputModel,
			InputUseRegions,
			InputThreshold,
			InputMinPercentage,
			InputMaxRectangles,
			InputPadding
		],
		outputs=[OutputVideo, OutputDownload],
	)

if __name__ == '__main__':
	os.makedirs(ModelDir, exist_ok=True)
	os.makedirs(TempDir, exist_ok=True)
	Logger.info('🚀 Starting Video Upscaler')
	Interface.launch(pwa=True)