import copy
import logging
import math

import torch


FLOAT_TENSOR_SCOPE_ALL = "all_float_tensors"
FLOAT_TENSOR_SCOPE_MAIN_AND_METADATA = "main_conditioning_and_float_metadata"
FLOAT_TENSOR_SCOPE_MAIN_ONLY = "main_conditioning_only"
FLOAT_TENSOR_SCOPE_METADATA_ONLY = "float_metadata_only"
FLOAT_TENSOR_SCOPE_CHOICES = [
	FLOAT_TENSOR_SCOPE_MAIN_AND_METADATA,
	FLOAT_TENSOR_SCOPE_MAIN_ONLY,
	FLOAT_TENSOR_SCOPE_METADATA_ONLY,
	FLOAT_TENSOR_SCOPE_ALL,
]

NON_FLOAT_BEHAVIOR_PRESERVE = "preserve"
NON_FLOAT_BEHAVIOR_ERROR = "error"
NON_FLOAT_BEHAVIOR_CHOICES = [
	NON_FLOAT_BEHAVIOR_PRESERVE,
	NON_FLOAT_BEHAVIOR_ERROR,
]

CURVE_LINEAR = "linear"
CURVE_COSINE = "cosine"
CURVE_SMOOTHSTEP = "smoothstep"
CURVE_EASE_IN = "ease_in"
CURVE_EASE_OUT = "ease_out"
CURVE_SIGMOID = "sigmoid"
CURVE_CHOICES = [
	CURVE_LINEAR,
	CURVE_COSINE,
	CURVE_SMOOTHSTEP,
	CURVE_EASE_IN,
	CURVE_EASE_OUT,
	CURVE_SIGMOID,
]

OUTSIDE_WINDOW_HOLD = "hold"
OUTSIDE_WINDOW_BASELINE = "baseline"
OUTSIDE_WINDOW_LINEAR_EXTRAPOLATE = "linear_extrapolate"
OUTSIDE_WINDOW_CHOICES = [
	OUTSIDE_WINDOW_HOLD,
	OUTSIDE_WINDOW_BASELINE,
	OUTSIDE_WINDOW_LINEAR_EXTRAPOLATE,
]

TIMESTEP_EPSILON = 1e-5


def _path_matches_metadata_key(path, metadata_keys):
	if not path:
		return False
	last_key = path[-1]
	return isinstance(last_key, str) and last_key in metadata_keys


def _is_main_conditioning_tensor_path(path):
	if len(path) == 1:
		return path[0] == 0
	return len(path) == 2 and isinstance(path[0], int) and path[1] == 0


def _should_multiply_float_tensor(path, tensor_scope, metadata_keys):
	is_main_conditioning = _is_main_conditioning_tensor_path(path)
	is_selected_metadata = _path_matches_metadata_key(path, metadata_keys)

	if tensor_scope == FLOAT_TENSOR_SCOPE_ALL:
		return True
	if tensor_scope == FLOAT_TENSOR_SCOPE_MAIN_ONLY:
		return is_main_conditioning
	if tensor_scope == FLOAT_TENSOR_SCOPE_METADATA_ONLY:
		return is_selected_metadata
	return is_main_conditioning or is_selected_metadata


def _multiply_structure(
	structure,
	multiplier,
	tensor_scope,
	non_float_behavior,
	metadata_keys,
	path=(),
	stats=None,
):
	if stats is None:
		stats = {
			"multiplied": 0,
			"preserved_non_float": 0,
			"preserved_float": 0,
		}

	if isinstance(structure, torch.Tensor):
		if structure.is_floating_point() or structure.is_complex():
			if _should_multiply_float_tensor(path, tensor_scope, metadata_keys):
				stats["multiplied"] += 1
				return structure * multiplier
			stats["preserved_float"] += 1
			return structure

		if non_float_behavior == NON_FLOAT_BEHAVIOR_ERROR:
			raise TypeError(
				"ConditioningMultiplyAdvanced refused to multiply non-floating tensor "
				f"at path {'.'.join(map(str, path)) or '<root>'} with dtype {structure.dtype}."
			)
		stats["preserved_non_float"] += 1
		return structure

	if isinstance(structure, list):
		return [
			_multiply_structure(
				item,
				multiplier,
				tensor_scope,
				non_float_behavior,
				metadata_keys,
				path + (index,),
				stats,
			)
			for index, item in enumerate(structure)
		]

	if isinstance(structure, tuple):
		return tuple(
			_multiply_structure(
				item,
				multiplier,
				tensor_scope,
				non_float_behavior,
				metadata_keys,
				path + (index,),
				stats,
			)
			for index, item in enumerate(structure)
		)

	if isinstance(structure, dict):
		return {
			key: _multiply_structure(
				value,
				multiplier,
				tensor_scope,
				non_float_behavior,
				metadata_keys,
				path + (key,),
				stats,
			)
			for key, value in structure.items()
		}

	return copy.deepcopy(structure)


def _parse_metadata_keys(metadata_keys):
	if not isinstance(metadata_keys, str):
		return set()
	return {
		key.strip()
		for key in metadata_keys.split(",")
		if key.strip()
	}


def _clamp_percent(value):
	return max(0.0, min(1.0, float(value)))


def _curve_weight(percent, curve):
	percent = _clamp_percent(percent)

	if curve == CURVE_LINEAR:
		return percent
	if curve == CURVE_SMOOTHSTEP:
		return percent * percent * (3.0 - 2.0 * percent)
	if curve == CURVE_COSINE:
		return 0.5 - 0.5 * math.cos(math.pi * percent)
	if curve == CURVE_EASE_IN:
		return percent * percent
	if curve == CURVE_EASE_OUT:
		return 1.0 - (1.0 - percent) * (1.0 - percent)
	if curve == CURVE_SIGMOID:
		return 1.0 / (1.0 + math.exp(-10.0 * (percent - 0.5)))

	raise ValueError(f"Unknown curve: {curve}")


def _lerp(start, end, weight):
	return (1.0 - weight) * start + weight * end


def _windowed_multiplier(
	percent,
	start_multiplier,
	end_multiplier,
	start_percent,
	end_percent,
	curve,
	outside_window,
):
	denominator = max(TIMESTEP_EPSILON, end_percent - start_percent)
	unbounded_percent = (percent - start_percent) / denominator

	if start_percent <= percent <= end_percent:
		return _lerp(start_multiplier, end_multiplier, _curve_weight(unbounded_percent, curve))

	if outside_window == OUTSIDE_WINDOW_BASELINE:
		return 1.0
	if outside_window == OUTSIDE_WINDOW_LINEAR_EXTRAPOLATE:
		return _lerp(start_multiplier, end_multiplier, unbounded_percent)
	if percent < start_percent:
		return start_multiplier
	return end_multiplier


def _copy_conditioning_entry_with_values(entry, values):
	if len(entry) < 2 or not isinstance(entry[1], dict):
		return copy.deepcopy(entry)

	metadata = entry[1].copy()
	metadata.update(values)
	return [entry[0], metadata]


def _get_conditioning_range(entry):
	if len(entry) < 2 or not isinstance(entry[1], dict):
		return 0.0, 1.0
	return (
		_clamp_percent(entry[1].get("start_percent", 0.0)),
		_clamp_percent(entry[1].get("end_percent", 1.0)),
	)


def _append_scaled_range(
	out,
	entry,
	range_start,
	range_end,
	multiplier,
	tensor_scope,
	non_float_behavior,
	metadata_keys,
	stats,
):
	if range_start >= range_end:
		return

	scaled_entry = _multiply_structure(
		entry,
		float(multiplier),
		tensor_scope,
		non_float_behavior,
		metadata_keys,
		stats=stats,
	)
	out.append(_copy_conditioning_entry_with_values(scaled_entry, {
		"start_percent": range_start,
		"end_percent": range_end,
	}))


def _append_segmented_range(
	out,
	entry,
	range_start,
	range_end,
	segments,
	start_multiplier,
	end_multiplier,
	start_percent,
	end_percent,
	curve,
	outside_window,
	tensor_scope,
	non_float_behavior,
	metadata_keys,
	stats,
):
	if range_start >= range_end:
		return

	segments = max(1, int(segments))
	span = range_end - range_start
	for segment_index in range(segments):
		segment_start = range_start + span * (segment_index / segments)
		segment_end = range_start + span * ((segment_index + 1) / segments)
		if segment_index < segments - 1:
			segment_end = max(segment_start, segment_end - TIMESTEP_EPSILON)

		midpoint = (segment_start + segment_end) * 0.5
		multiplier = _windowed_multiplier(
			midpoint,
			start_multiplier,
			end_multiplier,
			start_percent,
			end_percent,
			curve,
			outside_window,
		)
		_append_scaled_range(
			out,
			entry,
			segment_start,
			segment_end,
			multiplier,
			tensor_scope,
			non_float_behavior,
			metadata_keys,
			stats,
		)


def _multiply_conditioning_with_schedule(
	conditioning,
	start_multiplier,
	end_multiplier,
	start_percent,
	end_percent,
	curve,
	outside_window,
	segments,
	tensor_scope,
	non_float_behavior,
	metadata_keys,
	stats,
):
	start_percent = _clamp_percent(start_percent)
	end_percent = _clamp_percent(end_percent)
	segments = max(1, int(segments))

	if (
		abs(float(start_multiplier) - float(end_multiplier)) <= 1e-12
		and (
			outside_window == OUTSIDE_WINDOW_HOLD
			or (start_percent <= 0.0 and end_percent >= 1.0)
		)
	):
		return _multiply_structure(
			conditioning,
			float(start_multiplier),
			tensor_scope,
			non_float_behavior,
			metadata_keys,
			stats=stats,
		)

	if start_percent >= end_percent:
		logging.warning(
			"ConditioningMultiplyAdvanced: start_percent must be lower than end_percent; "
			"returning unchanged conditioning."
		)
		return conditioning

	out = []
	for entry in conditioning:
		conditioning_start, conditioning_end = _get_conditioning_range(entry)
		intersect_start = max(start_percent, conditioning_start)
		intersect_end = min(end_percent, conditioning_end)

		if conditioning_start < start_percent:
			before_end = min(conditioning_end, max(conditioning_start, start_percent - TIMESTEP_EPSILON))
			before_segments = segments if outside_window == OUTSIDE_WINDOW_LINEAR_EXTRAPOLATE else 1
			_append_segmented_range(
				out,
				entry,
				conditioning_start,
				before_end,
				before_segments,
				start_multiplier,
				end_multiplier,
				start_percent,
				end_percent,
				curve,
				outside_window,
				tensor_scope,
				non_float_behavior,
				metadata_keys,
				stats,
			)

		if intersect_start < intersect_end:
			_append_segmented_range(
				out,
				entry,
				intersect_start,
				intersect_end,
				segments,
				start_multiplier,
				end_multiplier,
				start_percent,
				end_percent,
				curve,
				outside_window,
				tensor_scope,
				non_float_behavior,
				metadata_keys,
				stats,
			)

		if end_percent < conditioning_end:
			after_start = max(conditioning_start, min(conditioning_end, end_percent + TIMESTEP_EPSILON))
			after_segments = segments if outside_window == OUTSIDE_WINDOW_LINEAR_EXTRAPOLATE else 1
			_append_segmented_range(
				out,
				entry,
				after_start,
				conditioning_end,
				after_segments,
				start_multiplier,
				end_multiplier,
				start_percent,
				end_percent,
				curve,
				outside_window,
				tensor_scope,
				non_float_behavior,
				metadata_keys,
				stats,
			)

	return out


class ConditioningMultiplyAdvanced:
	@classmethod
	def INPUT_TYPES(s):
		return {
			"required": {
				"conditioning": ("CONDITIONING", {
					"tooltip": "Conditioning payload to scale."
				}),
				"start_multiplier": ("FLOAT", {
					"default": 1.0,
					"min": -1000000000.0,
					"max": 1000000000.0,
					"step": 0.01,
					"round": False,
					"tooltip": "Multiplier used at start_percent and before the window when outside_window is hold."
				}),
				"end_multiplier": ("FLOAT", {
					"default": 1.0,
					"min": -1000000000.0,
					"max": 1000000000.0,
					"step": 0.01,
					"round": False,
					"tooltip": "Multiplier used at end_percent and after the window when outside_window is hold."
				}),
				"start_percent": ("FLOAT", {
					"default": 0.0,
					"min": 0.0,
					"max": 1.0,
					"step": 0.001,
					"tooltip": "Denoise progress where the multiplier transition window starts."
				}),
				"end_percent": ("FLOAT", {
					"default": 1.0,
					"min": 0.0,
					"max": 1.0,
					"step": 0.001,
					"tooltip": "Denoise progress where the multiplier transition window ends."
				}),
				"curve": (CURVE_CHOICES, {
					"default": CURVE_LINEAR,
					"tooltip": "Curve used to interpolate from start_multiplier to end_multiplier inside the timestep window."
				}),
				"outside_window": (OUTSIDE_WINDOW_CHOICES, {
					"default": OUTSIDE_WINDOW_HOLD,
					"tooltip": "hold uses start_multiplier before the window and end_multiplier after it. baseline uses 1.0 outside the window. linear_extrapolate continues a linear multiplier trend outside the window."
				}),
				"segments": ("INT", {
					"default": 16,
					"min": 1,
					"max": 256,
					"step": 1,
					"tooltip": "Number of conditioning ranges used to approximate the multiplier curve inside the window."
				}),
				"tensor_scope": (FLOAT_TENSOR_SCOPE_CHOICES, {
					"default": FLOAT_TENSOR_SCOPE_MAIN_AND_METADATA,
					"tooltip": "Controls which floating tensors are multiplied. The default scales the main conditioning tensor plus selected float metadata such as pooled_output and t5xxl_weights."
				}),
				"non_float_behavior": (NON_FLOAT_BEHAVIOR_CHOICES, {
					"default": NON_FLOAT_BEHAVIOR_PRESERVE,
					"tooltip": "preserve leaves integer/bool tensors unchanged. error raises when a non-floating tensor is encountered."
				}),
				"metadata_keys": ("STRING", {
					"default": "pooled_output,t5xxl_weights",
					"multiline": False,
					"tooltip": "Comma-separated metadata tensor keys to scale when tensor_scope includes metadata. Integer token ids such as t5xxl_ids should not be listed."
				}),
				"log_summary": ("BOOLEAN", {
					"default": False,
					"tooltip": "Print how many tensors were multiplied or preserved."
				}),
			}
		}

	RETURN_TYPES = ("CONDITIONING",)
	RETURN_NAMES = ("conditioning",)
	FUNCTION = "multiply"
	CATEGORY = "advanced/conditioning"
	DESCRIPTION = "Multiply selected floating conditioning tensors while preserving integer token-id tensors such as Anima t5xxl_ids."
	SEARCH_ALIASES = ["multiply conditioning", "scale conditioning", "conditioning strength", "prompt strength"]

	def multiply(
		self,
		conditioning,
		start_multiplier,
		end_multiplier,
		start_percent,
		end_percent,
		curve,
		outside_window,
		segments,
		tensor_scope=FLOAT_TENSOR_SCOPE_MAIN_AND_METADATA,
		non_float_behavior=NON_FLOAT_BEHAVIOR_PRESERVE,
		metadata_keys="pooled_output,t5xxl_weights",
		log_summary=False,
	):
		if tensor_scope not in FLOAT_TENSOR_SCOPE_CHOICES:
			tensor_scope = FLOAT_TENSOR_SCOPE_MAIN_AND_METADATA
		if non_float_behavior not in NON_FLOAT_BEHAVIOR_CHOICES:
			non_float_behavior = NON_FLOAT_BEHAVIOR_PRESERVE
		if curve not in CURVE_CHOICES:
			curve = CURVE_LINEAR
		if outside_window not in OUTSIDE_WINDOW_CHOICES:
			outside_window = OUTSIDE_WINDOW_HOLD

		stats = {
			"multiplied": 0,
			"preserved_non_float": 0,
			"preserved_float": 0,
		}
		out = _multiply_conditioning_with_schedule(
			conditioning,
			float(start_multiplier),
			float(end_multiplier),
			start_percent,
			end_percent,
			curve,
			outside_window,
			segments,
			tensor_scope,
			non_float_behavior,
			_parse_metadata_keys(metadata_keys),
			stats,
		)

		if log_summary:
			logging.info(
				"ConditioningMultiplyAdvanced: "
				f"multiplied={stats['multiplied']} "
				f"preserved_float={stats['preserved_float']} "
				f"preserved_non_float={stats['preserved_non_float']} "
				f"scope={tensor_scope} "
				f"start_multiplier={float(start_multiplier):.4g} "
				f"end_multiplier={float(end_multiplier):.4g} "
				f"start_percent={_clamp_percent(start_percent):.3f} "
				f"end_percent={_clamp_percent(end_percent):.3f} "
				f"curve={curve} "
				f"outside_window={outside_window} "
				f"segments={max(1, int(segments))}"
			)

		return (out,)


NODE_CLASS_MAPPINGS = {
	"ConditioningMultiplyAdvanced": ConditioningMultiplyAdvanced,
}

NODE_DISPLAY_NAME_MAPPINGS = {
	"ConditioningMultiplyAdvanced": "Conditioning Multiply Advanced",
}

__all__ = [
	"NODE_CLASS_MAPPINGS",
	"NODE_DISPLAY_NAME_MAPPINGS",
]