from itertools import repeat
import logging
from pathlib import Path
import re
import shutil
import torch
import glm
import csv
import numpy as np
from typing import List, Union
from torch.types import Number
from .constants import *
from .device import *


gvec_type = [glm.dvec1, glm.dvec2, glm.dvec3, glm.dvec4]
gmat_type = [[glm.dmat2, glm.dmat2x3, glm.dmat2x4],
             [glm.dmat3x2, glm.dmat3, glm.dmat3x4],
             [glm.dmat4x2, glm.dmat4x3, glm.dmat4]]


def smooth_step(x0, x1, x):
    y = torch.clamp((x - x0) / (x1 - x0), 0, 1)
    return y * y * (3 - 2 * y)


def torch2np(input: torch.Tensor) -> np.ndarray:
    return input.cpu().detach().numpy()


def torch2glm(input):
    input = input.squeeze()
    size = input.size()
    if len(size) == 1:
        if size[0] <= 0 or size[0] > 4:
            raise ValueError
        return gvec_type[size[0] - 1](torch2np(input))
    if len(size) == 2:
        if size[0] <= 1 or size[0] > 4 or size[1] <= 1 or size[1] > 4:
            raise ValueError
        return gmat_type[size[1] - 2][size[0] - 2](torch2np(input))
    raise ValueError


def glm2torch(val) -> torch.Tensor:
    return torch.from_numpy(np.array(val))


def meshgrid(*size: int, normalize: bool = False, swap_dim: bool = False) -> torch.Tensor:
    """
    Generate a mesh grid

    :param *size: grid size (rows, columns)
    :param normalize: return coords in normalized space? defaults to False
    :param swap_dim: if True, return coords in (y, x) order, defaults to False
    :return: rows x columns x 2 tensor
    """
    if len(size) == 1:
        size = (size[0], size[0])
    y, x = torch.meshgrid(torch.arange(size[0]), torch.arange(size[1]), indexing='ij')
    if normalize:
        x.div_(size[1] - 1.)
        y.div_(size[0] - 1.)
    return torch.stack([y, x], 2) if swap_dim else torch.stack([x, y], 2)


def get_angle(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    angle = -torch.atan(x / y) - (y < 0) * PI + 0.5 * PI
    return angle


def broadcast_cat(input: torch.Tensor,
                  s: Union[Number, List[Number], torch.Tensor],
                  dim=-1,
                  append: bool = True) -> torch.Tensor:
    """
    Concatenate a tensor with a scalar along last dimension

    :param input `Tensor(..., N)`: input tensor
    :param s: scalar
    :param append: append or prepend the scalar to input tensor
    :return: `Tensor(..., N+1)`
    """
    if dim != -1:
        raise NotImplementedError('currently only support the last dimension')
    if isinstance(s, torch.Tensor):
        x = s
    elif isinstance(s, list):
        x = torch.tensor(s, dtype=input.dtype, device=input.device)
    else:
        x = torch.tensor([s], dtype=input.dtype, device=input.device)
    expand_shape = list(input.size())
    expand_shape[dim] = -1
    x = x.expand(expand_shape)
    return torch.cat([input, x] if append else [x, input], dim)


def save_2d_tensor(path, x):
    with open(path, 'w', encoding='utf-8', newline='') as f:
        csv_writer = csv.writer(f)
        for i in range(x.shape[0]):
            csv_writer.writerow(x[i])


def view_like(input: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
    """
    Reshape input to be the same size as ref except the last dimension

    :param input `Tensor(..., C)`: input tensor
    :param ref `Tensor(B.., *): reference tensor
    :return `Tensor(B.., C)`: reshaped tensor
    """
    out_shape = list(ref.size())
    out_shape[-1] = -1
    return input.view(out_shape)


def format_time(seconds):
    days = int(seconds / 3600 / 24)
    seconds = seconds - days * 3600 * 24
    hours = int(seconds / 3600)
    seconds = seconds - hours * 3600
    minutes = int(seconds / 60)
    seconds = seconds - minutes * 60
    seconds_final = int(seconds)
    seconds = seconds - seconds_final
    millis = int(seconds * 1000)

    if days > 0:
        output = f"{days}D{hours:0>2d}h{minutes:0>2d}m"
    elif hours > 0:
        output = f"{hours:0>2d}h{minutes:0>2d}m{seconds_final:0>2d}s"
    elif minutes > 0:
        output = f"{minutes:0>2d}m{seconds_final:0>2d}s"
    elif seconds_final > 0:
        output = f"{seconds_final:0>2d}s{millis:0>3d}ms"
    elif millis > 0:
        output = f"{millis:0>3d}ms"
    else:
        output = '0ms'
    return output


def print_and_log(s):
    print(s)
    logging.info(s)


def masked_scatter(mask: torch.Tensor, value: torch.Tensor, initial: Union[torch.Tensor, Number] = 0):
    """
    Extend PyTorch's built-in `masked_scatter` function

    :param mask `Tensor(M...)`: the boolean mask
    :param value `Tensor(N, D...)`: the value to fill in with, should have at least as many elements 
                                    as the number of ones in `mask`
    :param destination `Tensor(M..., D...)`: (optional) the destination tensor to fill,
                                             if not specified, a new tensor filled with 
                                             `empty_value` will be created and used as destination
    :param empty_value `Number`: the initial elements in the newly created destination tensor, 
                                 defaults to 0
    :return `Tensor(M..., D...)`: the destination tensor after filled
    """
    M_ = mask.size()
    D_ = value.size()[1:]
    if not isinstance(initial, torch.Tensor):
        initial = value.new_full([*M_, *D_], initial)
    return initial.masked_scatter(mask.reshape(*M_, *repeat(1, len(D_))), value)


def list_epochs(dir: Path, pattern: str) -> List[int]:
    prefix = pattern.split("*")[0]
    epoch_list = [int(str(path.stem)[len(prefix):]) for path in dir.glob(pattern)]
    epoch_list.sort()
    return epoch_list


def rename_seqs_with_offset(dir: Path, file_pattern: str, offset: int):
    start, end = re.search(r'%0\dd', file_pattern).span()
    prefix, suffix = start, len(file_pattern) - end

    seqs = [
        int(path.name[prefix:-suffix])
        for path in dir.glob(re.sub(r'%0\dd', "*", file_pattern))
    ]
    seqs.sort(reverse=offset > 0)
    for i in seqs:
        (dir / (file_pattern % i)).rename(dir / (file_pattern % (i + offset)))