import re
from itertools import product
from typing import Callable, List, Dict, Any, Union, Tuple, cast
import torch
import model_management # type: ignore
import comfy.samplers
from nodes import common_ksampler
from comfy.sd import ModelPatcher
from .model.iter import iterize_model, CondForModels

re_int = re.compile(r"\s*([+-]?\s*\d+)\s*")
re_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*")
re_range = re.compile(r"\s*([+-]?\s*\d+)\s*-\s*([+-]?\s*\d+)(?:\s*\(([+-]\d+)\s*\))?\s*")
re_range_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+(?:.\d*)?)(?:\s*\(([+-]\d+(?:.\d*)?)\s*\))?\s*")

def frange(start, end, step):
    x = float(start)
    end = float(end)
    step = float(step)
    while x < end:
        yield x
        x += step

def get_noise(seeds: List[int], latent_image: torch.Tensor, disable_noise: bool):
    noises: List[torch.Tensor] = []
    latents: List[torch.Tensor] = []
    
    if latent_image.dim() == 3:
        latent_image = latent_image.unsqueeze(0) # add batch dim
    
    if disable_noise:
        noise_ = torch.zeros([len(seeds)]+list(latent_image.size())[-3:], dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
        noises.append(noise_)
        latents.extend([latent_image] * (len(seeds) // latent_image.shape[0]))
    else:
        for s in seeds:
            noise_ = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=torch.manual_seed(s), device="cpu")
            noises.append(noise_)
            latents.append(latent_image)
    
    return torch.cat(noises), torch.cat(latents)

def get_cfg(noises: torch.Tensor, latent_image: torch.Tensor, cfgs: List[float]):
    # batch_size = noises.shape[0] * len(cfgs)
    ns = [noises] * len(cfgs)
    lat = [latent_image] * len(cfgs)
    cf = torch.FloatTensor(cfgs * noises.shape[0])
    return torch.cat(ns), torch.cat(lat), cf[...,None,None,None]

def process_cond(
    conds: List[List[Union[torch.Tensor,dict]]],
    control_nets: Union[list,None],
    noise_shape0: int,
    device
):
    conds_copy = []
    for p in conds:
        t: torch.Tensor = p[0] # type: ignore
        if t.shape[0] < noise_shape0:
            t = torch.cat([t] * noise_shape0)
        t = t.to(device)
        if control_nets is not None and 'control' in p[1]:
            control_nets += [p[1]['control']] # type: ignore
        conds_copy += [[t] + p[1:]]
    return conds_copy

def process_cond_for_models(
    conds: List[List[Union[torch.Tensor,CondForModels,dict]]],
    control_nets: list,
    noise_shape0: int,
    device
):
    assert (
        all(isinstance(p[0], CondForModels) for p in conds) 
        or not any(isinstance(p[0], CondForModels) for p in conds)
    )
    
    if isinstance(conds[0][0], CondForModels):
        sizes = set( len(cast(CondForModels, p[0]).ex) for p in conds )
        assert len(sizes) == 1, f'number of conditions: {sizes}'
        size = sizes.pop()
        
        #
        # conds
        #  + [ CondForModels, dictA ]
        #  |       .ex + condA for model1
        #  |           + condA for model2
        #  |           ...
        #  |           L condA for model{size}
        #  + [ CondForModels, dictB ]
        #  |       .ex + condB for model1
        #  |           + condB for model2
        #  |           ...
        #  |           L condB for model{size}
        #  ...
        # 
        # vvv
        # 
        # conds
        #  + [ [ condA_for_model1, dictA ], [ condB_for_model1, dictB ], ... ]
        #  + [ [ condA_for_model2, dictA ], [ condB_for_model2, dictB ], ... ]
        #  ...
        # 
        
        result = []
        for model_index in range(size):
            cs = []
            for cond_for_models in conds:
                c: CondForModels = cond_for_models[0] # type: ignore
                rest = cond_for_models[1:]
                cond = c.ex[model_index]
                cs.append([cond, *rest])
            result.append(process_cond(
                cs,
                control_nets if model_index == 0 else None,
                noise_shape0,
                device
            ))
        return result
    
    else:
        return [ process_cond(conds, control_nets, noise_shape0, device) ]
    

def common_ksampler_xyz(
    model: ModelPatcher,
    seed: Union[int,List[int]],
    steps: Union[int,List[int]],
    cfg: Union[float,List[float]],
    sampler_name: Union[str,List[str]],
    scheduler: Union[str,List[str]],
    positive,
    negative,
    latent,
    denoise=1.0,
    disable_noise=False,
    start_step=None,
    last_step=None,
    force_full_denoise=False
):
    latent_image = latent["samples"]
    noise_mask = None
    device = model_management.get_torch_device()

    if not isinstance(seed, list):
        seed = [seed]
    
    if not isinstance(steps, list):
        steps = [steps]
    
    if not isinstance(cfg, list):
        cfg = [cfg]
    
    if not isinstance(sampler_name, list):
        sampler_name = [sampler_name]
    
    if not isinstance(scheduler, list):
        scheduler = [scheduler]
    
    noise, latent_image = get_noise(seed, latent_image, disable_noise)
    noise, latent_image, cfg_ = get_cfg(noise, latent_image, cfg)
    
    if "noise_mask" in latent:
        noise_mask = latent['noise_mask']
        noise_mask = torch.nn.functional.interpolate(noise_mask[None,None,], size=(noise.shape[2], noise.shape[3]), mode="bilinear")
        noise_mask = noise_mask.round()
        noise_mask = torch.cat([noise_mask] * noise.shape[1], dim=1)
        noise_mask = torch.cat([noise_mask] * noise.shape[0])
        noise_mask = noise_mask.to(device)

    noise = noise.to(device)
    latent_image = latent_image.to(device)
    cfg_ = cfg_.to(device)

    control_nets = []
    positive_copies = process_cond_for_models(positive, control_nets, noise.shape[0], device)
    negative_copies = process_cond_for_models(negative, control_nets, noise.shape[0], device)

    control_net_models = []
    for x in control_nets:
        control_net_models += x.get_control_models()
    model_management.load_controlnet_gpu(control_net_models)

    samplers: List[Dict[str,Any]] = []
    for (model_index, model_fn), sampler_name_, scheduler_, steps_ in product(enumerate(iterize_model(model)), sampler_name, scheduler, steps):
        if sampler_name_ not in comfy.samplers.KSampler.SAMPLERS:
            raise ValueError(f'unknown sampler name: {sampler_name_}')
        if scheduler_ not in comfy.samplers.KSampler.SCHEDULERS:
            raise ValueError(f'unknown scheduler name: {scheduler_}')
        samplers.append(dict(
            model_index=model_index,
            model=model_fn,
            steps=steps_,
            device=device,
            sampler=sampler_name_,
            scheduler=scheduler_,
            denoise=denoise,
        ))
    
    all_samples: List[torch.Tensor] = []
    for sampler_args in samplers:
        model_ = sampler_args['model']()
        model_management.load_model_gpu(model_)
        sampler_args['model'] = model_.model
        
        model_index = sampler_args.pop('model_index')
        positive_copy = positive_copies[model_index]
        negative_copy = negative_copies[model_index]
        
        sampler = comfy.samplers.KSampler(**sampler_args)
        print(f'XYZ sampler=model@{model_index}/{sampler.sampler}/{sampler.scheduler} {sampler.steps}steps')
        
        samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg_, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask)
        samples = samples.cpu()
        all_samples.append(samples)
    for c in control_nets:
        c.cleanup()

    out = latent.copy()
    out["samples"] = torch.cat(all_samples)
    return (out, )


class KSamplerSetting:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'model': ('MODEL',),
                'seed': ('INT', {'default': 0, 'min': 0, 'max': 0xffffffffffffffff}),
                'steps': ('INT', {'default': 20, 'min': 1, 'max': 10000}),
                'cfg': ('FLOAT', {'default': 8.0, 'min': 0.0, 'max': 100.0}),
                'sampler_name': (comfy.samplers.KSampler.SAMPLERS, ),
                'scheduler': (comfy.samplers.KSampler.SCHEDULERS, ),
                'positive': ('CONDITIONING', ),
                'negative': ('CONDITIONING', ),
                'latent_image': ('LATENT', ),
                'denoise': ('FLOAT', {'default': 1.0, 'min': 0.0, 'max': 1.0, 'step': 0.01}),
            }
        }

    RETURN_TYPES = ('DICT',)
    
    FUNCTION = 'sample'

    CATEGORY = 'sampling'

    def sample(self, **kwargs):
        return kwargs,


class KSamplerOverrided:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'setting': ('DICT',),
            },
            'optional': {
                'model': ('MODEL',),
                'seed': ('Integer', {'default': 0, 'min': 0, 'max': 0xffffffffffffffff}),
                'steps': ('Integer', {'default': 20, 'min': 1, 'max': 10000}),
                'cfg': ('Float', {'default': 8.0, 'min': 0.0, 'max': 100.0}),
                'sampler_name': ('SamplerName',),
                'scheduler': ('SchedulerName', ),
                'positive': ('CONDITIONING', ),
                'negative': ('CONDITIONING', ),
                'latent_image': ('LATENT', ),
                'denoise': ('Float', {'default': 1.0, 'min': 0.0, 'max': 1.0, 'step': 0.01}),
            }
        }

    RETURN_TYPES = ('LATENT',)
    FUNCTION = 'sample'

    CATEGORY = 'sampling'

    def sample(self, setting: dict, **kwargs):
        if 'latent_image' in setting:
            setting['latent'] = setting['latent_image']
            del setting['latent_image']
        
        setting.update(kwargs)
        
        return common_ksampler(**setting)

class KSamplerXYZ:
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'setting': ('DICT',),
            },
            'optional': {
                'model': ('MODEL',),
                'seed': ('STRING', { 'multiline': True, 'default': '' }),
                'steps': ('STRING', { 'multiline': True, 'default': '' }),
                'cfg': ('STRING', { 'multiline': True, 'default': '' }),
                'sampler_name': ('STRING', { 'multiline': True, 'default': '' }),
                'scheduler': ('STRING', { 'multiline': True, 'default': '' }),
            }
        }

    RETURN_TYPES = ('LATENT',)
    FUNCTION = 'sample'

    CATEGORY = 'sampling'

    def sample(self, setting: dict, **kwargs):
        if 'latent_image' in setting:
            setting['latent'] = setting['latent_image']
            del setting['latent_image']
        
        # ignore empty string
        kwargs = { k: v for k, v in kwargs.items() if not isinstance(v, str) or len(v) != 0 }
        
        setting = { **setting, **kwargs }
        
        if isinstance(setting.get('seed', None), str):
            setting['seed'] = self.parse(setting['seed'], self.parse_int)
            
        if isinstance(setting.get('steps', None), str):
            setting['steps'] = self.parse(setting['steps'], self.parse_int)
        
        if isinstance(setting.get('cfg', None), str):
            setting['cfg'] = self.parse(setting['cfg'], self.parse_float)
        
        if isinstance(setting.get('sampler_name', None), str):
            setting['sampler_name'] = self.parse(setting['sampler_name'], None)
            if len(setting['sampler_name']) == 1:
                setting['sampler_name'] = setting['sampler_name'][0]
        
        if isinstance(setting.get('scheduler', None), str):
            setting['scheduler'] = self.parse(setting['scheduler'], None)
            if len(setting['scheduler']) == 1:
                setting['scheduler'] = setting['scheduler'][0]
        
        for k, v in setting.items():
            if k in kwargs and isinstance(v, (list, tuple)):
                print(f'XYZ {k}: {v}')
        
        return common_ksampler_xyz(**setting) # type: ignore
    
    def parse(self, input: str, cont: Union[Callable[[str],Any],None]):
        vs = [ x.strip() for x in input.split(',') ]
        if cont is not None:
            vs = [cont(v) for v in vs ]
        return vs
    
    def parse_int(self, input: str):
        m = re_int.fullmatch(input)
        if m is not None:
            return int(m.group(1))
        
        m = re_range.fullmatch(input)
        if m is None:
            raise ValueError(f'failed to process: {input}')
        
        start, end, step = m.group(1), m.group(2), m.group(3)
        if step is None:
            step = 1
        
        return list(range(int(start), int(end), int(step)))
    
    def parse_float(self, input: str):
        m = re_float.fullmatch(input)
        if m is not None:
            return float(m.group(1))
        
        m = re_range_float.fullmatch(input)
        if m is None:
            raise ValueError(f'failed to process: {input}')
        
        start, end, step = m.group(1), m.group(2), m.group(3)
        if step is None:
            step = 1.0
        
        return list(frange(float(start), float(end), float(step)))