import os
import sys
import time
import math
import numpy as np
from scipy.ndimage import gaussian_filter
from scipy.interpolate import CubicSpline as CubSpline
from scipy.special import comb
import scipy
from imageio import imread

import torch
import torch.nn.functional as F

# from perlin import PerlinNoiseFactory as Perlin
# noise = Perlin(1)

# def latent_noise(t, dim, noise_step=78564.543):
    # latent = np.zeros((1, dim))
    # for i in range(dim):
        # latent[0][i] = noise(t + i * noise_step)
    # return latent

def load_latents(npy_file):
    key_latents = np.load(npy_file)
    try:
        key_latents = key_latents[key_latents.files[0]]
    except:
        pass
    idx_file = os.path.splitext(npy_file)[0] + '.txt'
    if os.path.exists(idx_file): 
        with open(idx_file) as f:
            lat_idx = f.readline()
            lat_idx = [int(l.strip()) for l in lat_idx.split(',') if '\n' not in l and len(l.strip())>0]
        key_latents = [key_latents[i] for i in lat_idx]
    return np.asarray(key_latents)

# = = = = = = = = = = = = = = = = = = = = = = = = = = = 

def get_z(shape, seed=None, uniform=False):
    if seed is None:
        seed = np.random.seed(int((time.time()%1) * 9999))
    rnd = np.random.RandomState(seed)
    if uniform:
        return rnd.uniform(0., 1., shape)
    else:
        return rnd.randn(*shape) # *x unpacks tuple/list to sequence

def smoothstep(x, NN=1., xmin=0., xmax=1.):
    N = math.ceil(NN)
    x = np.clip((x - xmin) / (xmax - xmin), 0, 1)
    result = 0
    for n in range(0, N+1):
         result += scipy.special.comb(N+n, n) * scipy.special.comb(2*N+1, N-n) * (-x)**n
    result *= x**(N+1)
    if NN != N: result = (x + result) / 2
    return result

def lerp(z1, z2, num_steps, smooth=0.): 
    vectors = []
    xs = [step / (num_steps - 1) for step in range(num_steps)]
    if smooth > 0: xs = [smoothstep(x, smooth) for x in xs]
    for x in xs:
        interpol = z1 + (z2 - z1) * x
        vectors.append(interpol)
    return np.array(vectors)

# interpolate on hypersphere
def slerp(z1, z2, num_steps, smooth=0.):
    z1_norm = np.linalg.norm(z1)
    z2_norm = np.linalg.norm(z2)
    z2_normal = z2 * (z1_norm / z2_norm)
    vectors = []
    xs = [step / (num_steps - 1) for step in range(num_steps)]
    if smooth > 0: xs = [smoothstep(x, smooth) for x in xs]
    for x in xs:
        interplain = z1 + (z2 - z1) * x
        interp = z1 + (z2_normal - z1) * x
        interp_norm = np.linalg.norm(interp)
        interpol_normal = interplain * (z1_norm / interp_norm)
        # interpol_normal = interp * (z1_norm / interp_norm)
        vectors.append(interpol_normal)
    return np.array(vectors)

def cublerp(points, steps, fstep):
    keys = np.array([i*fstep for i in range(steps)] + [steps*fstep])
    points = np.concatenate((points, np.expand_dims(points[0], 0)))
    cspline = CubSpline(keys, points)
    return cspline(range(steps*fstep+1))

# = = = = = = = = = = = = = = = = = = = = = = = = = = = 
    
def latent_anima(shape, frames, transit, key_latents=None, smooth=0.5, cubic=False, gauss=False, seed=None, verbose=True):
    if key_latents is None:
        transit = int(max(1, min(frames//4, transit)))
    steps = max(1, int(frames // transit))
    log = ' timeline: %d steps by %d' % (steps, transit)

    getlat = lambda : get_z(shape, seed=seed)
    
    # make key points
    if key_latents is None:
        key_latents = np.array([getlat() for i in range(steps)])

    latents = np.expand_dims(key_latents[0], 0)
    
    # populate lerp between key points
    if transit == 1:
        latents = key_latents
    else:
        if cubic:
            latents = cublerp(key_latents, steps, transit)
            log += ', cubic'
        else:
            for i in range(steps):
                zA = key_latents[i]
                zB = key_latents[(i+1) % steps]
                interps_z = slerp(zA, zB, transit, smooth=smooth)
                latents = np.concatenate((latents, interps_z))
    latents = np.array(latents)
    
    if gauss:
        lats_post = gaussian_filter(latents, [transit, 0, 0], mode="wrap")
        lats_post = (lats_post / np.linalg.norm(lats_post, axis=-1, keepdims=True)) * math.sqrt(np.prod(shape))
        log += ', gauss'
        latents = lats_post
        
    if verbose: print(log)
    if latents.shape[0] > frames: # extra frame
        latents = latents[1:]
    return latents
    
# = = = = = = = = = = = = = = = = = = = = = = = = = = = 
    
def multimask(x, size, latmask=None, countHW=[1,1], delta=0.):
    Hx, Wx = countHW
    bcount = x.shape[0]

    if max(countHW) > 1:
        W = x.shape[3] # width
        H = x.shape[2] # height
        if Wx > 1:
            stripe_mask = []
            for i in range(Wx):
                ch_mask = peak_roll(W, Wx, i, delta).unsqueeze(0).unsqueeze(0) # [1,1,w] th
                ch_mask = ch_mask.repeat(1,H,1) # [1,h,w]
                stripe_mask.append(ch_mask)
            maskW = torch.cat(stripe_mask, 0).unsqueeze(1) # [x,1,h,w]
        else: maskW = [1]
        if Hx > 1:
            stripe_mask = []
            for i in range(Hx):
                ch_mask = peak_roll(H, Hx, i, delta).unsqueeze(1).unsqueeze(0) # [1,h,1] th
                ch_mask = ch_mask.repeat(1,1,W) # [1,h,w]
                stripe_mask.append(ch_mask)
            maskH = torch.cat(stripe_mask, 0).unsqueeze(1) # [y,1,h,w]
        else: maskH = [1]

        mask = []
        for i in range(Wx):
            for j in range(Hx):
                mask.append(maskW[i] * maskH[j])
        mask = torch.cat(mask, 0).unsqueeze(1) # [xy,1,h,w]
        mask = mask.to(x.device)
        x = torch.sum(x[:Hx*Wx] * mask, 0, keepdim=True)

    elif latmask is not None:
        if len(latmask.shape) < 4:
            latmask = latmask.unsqueeze(1) # [b,1,h,w]
        lms = latmask.shape
        if list(lms[2:]) != list(size) and np.prod(lms) > 1:
            latmask = F.interpolate(latmask, size) # , mode='nearest'
        latmask = latmask.type(x.dtype)
        x = torch.sum(x[:lms[0]] * latmask, 0, keepdim=True)
    else:
        return x

    x = x.repeat(bcount,1,1,1)
    return x # [b,f,h,w]

def peak_roll(width, count, num, delta):
    step = width // count
    if width > step*2:
        fill_range = torch.zeros([width-step*2])
        full_ax = torch.cat((peak(step, delta), fill_range), 0)
    else:
        full_ax = peak(step, delta)[:width]
    if num == 0: 
        shift = max(width - (step//2), 0.) # must be positive!
    else:
        shift = step*num - (step//2)
    full_ax = torch.roll(full_ax, shift, 0)
    return full_ax # [width,]

def peak(steps, delta):
    x = torch.linspace(0.-delta, 1.+ delta, steps)
    x_rev = torch.flip(x,[0])
    x = torch.cat((x, x_rev), 0)
    x = torch.clip(x, 0., 1.)
    return x # [steps*2,]

# = = = = = = = = = = = = = = = = = = = = = = = = = = = 
    
def ups2d(x, factor=2):
    assert isinstance(factor, int) and factor >= 1
    if factor == 1: return x
    s = x.shape
    x = x.reshape(-1, s[1], s[2], 1, s[3], 1)
    x = x.repeat(1, 1, 1, factor, 1, factor)
    x = x.reshape(-1, s[1], s[2] * factor, s[3] * factor)
    return x

# Tiles an array around two points, allowing for pad lengths greater than the input length
# NB: if symm=True, every second tile is mirrored = messed up in GAN
# adapted from https://discuss.pytorch.org/t/symmetric-padding/19866/3
def tile_pad(xt, padding, symm=True):
    h, w = xt.shape[-2:]
    left, right, top, bottom = padding
 
    def tile(x, minx, maxx, symm=True):
        rng = maxx - minx
        if symm is True: # triangular reflection
            double_rng = 2*rng
            mod = np.fmod(x - minx, double_rng)
            normed_mod = np.where(mod < 0, mod+double_rng, mod)
            out = np.where(normed_mod >= rng, double_rng - normed_mod, normed_mod) + minx
        else: # repeating tiles
            mod = np.remainder(x - minx, rng)
            out = mod + minx
        return np.array(out, dtype=x.dtype)

    x_idx = np.arange(-left, w+right)
    y_idx = np.arange(-top, h+bottom)
    x_pad = tile(x_idx, -0.5, w-0.5, symm)
    y_pad = tile(y_idx, -0.5, h-0.5, symm)
    xx, yy = np.meshgrid(x_pad, y_pad)
    return xt[..., yy, xx]

def pad_up_to(x, size, type='centr'):
    sh = x.shape[2:][::-1]
    if list(x.shape[2:]) == list(size): return x
    padding = []
    for i, s in enumerate(size[::-1]):
        if 'side' in type.lower():
            padding = padding + [0, s-sh[i]]
        else: # centr
            p0 = (s-sh[i]) // 2
            p1 = s-sh[i] - p0
            padding = padding + [p0,p1]
    y = tile_pad(x, padding, symm = 'symm' in type.lower())
    # if 'symm' in type.lower():
        # y = tile_pad(x, padding, symm=True)
    # else:
        # y = F.pad(x, padding, 'circular')
    return y

# scale_type may include pad, side, symm
def fix_size(x, size, scale_type='centr'): 
    if not len(x.shape) == 4:
        raise Exception(" Wrong data rank, shape:", x.shape)
    if x.shape[2:] == size:
        return x
    if (x.shape[2]*2, x.shape[3]*2) == size:
        return ups2d(x)

    if scale_type.lower() == 'fit':
        return F.interpolate(x, size, mode='nearest') # , align_corners=True
    elif 'pad' in scale_type.lower():
        pass
    else: # proportional scale to smaller side, then pad to bigger side
        sh0 = x.shape[2:]
        upsc = np.min(size) / np.min(sh0)
        new_size = [int(sh0[i]*upsc) for i in [0,1]]
        x = F.interpolate(x, new_size, mode='nearest') # , align_corners=True

    x = pad_up_to(x, size, scale_type)
    return x

# Make list of odd sizes for upsampling to arbitrary resolution
def hw_scales(size, base, n, keep_first_layers=None, verbose=False):
    if isinstance(base, int): base = (base, base)
    start_res = [int(b * 2 ** (-n)) for b in base]
    
    start_res[0] = int(start_res[0] * size[0] // base[0])
    start_res[1] = int(start_res[1] * size[1] // base[1])

    hw_list = []
    
    if base[0] != base[1] and verbose is True:
        print(' size', size, 'base', base, 'start_res', start_res, 'n', n)
    if keep_first_layers is not None and keep_first_layers > 0:
        for i in range(keep_first_layers):
            hw_list.append(start_res)
            start_res = [x*2 for x in start_res]
            n -= 1
            
    ch = (size[0] / start_res[0]) ** (1/n)
    cw = (size[1] / start_res[1]) ** (1/n)
    for i in range(n):
        h = math.floor(start_res[0] * ch**i)
        w = math.floor(start_res[1] * cw**i)
        hw_list.append((h,w))

    hw_list.append(size)
    return hw_list

def calc_res(shape):
    base0 = 2**int(np.log2(shape[0]))
    base1 = 2**int(np.log2(shape[1]))
    base = min(base0, base1)
    min_res = min(shape[0], shape[1])
    
    def int_log2(xs, base):
        return [x * 2**(2-int(np.log2(base))) % 1 == 0 for x in xs]
    if min_res != base or max(*shape) / min(*shape) >= 2:
        if np.log2(base) < 10 and all(int_log2(shape, base*2)):
            base = base * 2

    return base # , [shape[0]/base, shape[1]/base]

def calc_init_res(shape, resolution=None):
    if len(shape) == 1:
        shape = [shape[0], shape[0], 1]
    elif len(shape) == 2:
        shape = [*shape, 1]
    size = shape[:2] if shape[2] < min(*shape[:2]) else shape[1:] # fewer colors than pixels
    if resolution is None:
        resolution = calc_res(size)
    res_log2 = int(np.log2(resolution))
    init_res = [int(s * 2**(2-res_log2)) for s in size]
    return init_res, resolution, res_log2

def basename(file):
    return os.path.splitext(os.path.basename(file))[0]

def file_list(path, ext=None, subdir=None):
    if subdir is True:
        files = [os.path.join(dp, f) for dp, dn, fn in os.walk(path) for f in fn]
    else:
        files = [os.path.join(path, f) for f in os.listdir(path)]
    if ext is not None: 
        if isinstance(ext, list):
            files = [f for f in files if os.path.splitext(f.lower())[1][1:] in ext]
        elif isinstance(ext, str):
            files = [f for f in files if f.endswith(ext)]
        else:
            print(' Unknown extension/type for file list!')
    return sorted([f for f in files if os.path.isfile(f)])

def dir_list(in_dir):
    dirs = [os.path.join(in_dir, x) for x in os.listdir(in_dir)]
    return sorted([f for f in dirs if os.path.isdir(f)])

def img_list(path, subdir=None):
    if subdir is True:
        files = [os.path.join(dp, f) for dp, dn, fn in os.walk(path) for f in fn]
    else:
        files = [os.path.join(path, f) for f in os.listdir(path)]
    files = [f for f in files if os.path.splitext(f.lower())[1][1:] in ['jpg', 'jpeg', 'png', 'ppm', 'tif']]
    return sorted([f for f in files if os.path.isfile(f)])

def img_read(path):
    img = imread(path)
    # 8bit to 256bit
    if (img.ndim == 2) or (img.shape[2] == 1):
        img = np.dstack((img,img,img))
    # rgba to rgb 
    if img.shape[2] == 4:
        img = img[:,:,:3]
    return img