import torch
import math
from pathlib import Path
import time


def euler_rotation(rx=0, ry=0, rz=0, backwards=False):
    Ms = []
    if rz != 0:
        cosz = math.cos(rz)
        sinz = math.sin(rz)
        Ms.append(torch.tensor(
            [[cosz, -sinz, 0],
             [sinz, cosz, 0],
             [0, 0, 1]]))
    if ry != 0:
        cosy = math.cos(ry)
        siny = math.sin(ry)
        Ms.append(torch.tensor(
            [[cosy, 0, siny],
             [0, 1, 0],
             [-siny, 0, cosy]]))
    if rx != 0:
        cosx = math.cos(rx)
        sinx = math.sin(rx)
        Ms.append(torch.tensor(
            [[1, 0, 0],
             [0, cosx, -sinx],
             [0, sinx, cosx]]))

    rotation = torch.eye(3)
    if backwards and len(Ms) > 0:
        Ms = Ms[::-1]

    for mat in Ms[::-1]:
        rotation = torch.matmul(rotation, mat)
    return rotation


def export(file, vs, faces, vn=None, color=None):
    with open(file, 'w+') as f:
        for vi, v in enumerate(vs):
            if color is None:
                f.write("v %f %f %f\n" % (v[0], v[1], v[2]))
            else:
                f.write("v %f %f %f %f %f %f\n" % (v[0], v[1], v[2], color[vi][0], color[vi][1], color[vi][2]))
            if vn is not None:
                f.write("vn %f %f %f\n" % (vn[vi, 0], vn[vi, 1], vn[vi, 2]))
        for face in faces:
            f.write("f %d %d %d\n" % (face[0] + 1, face[1] + 1, face[2] + 1))


def xyz2tensor(txt, append_normals=False):
    pts = []
    for line in txt.split('\n'):
        line = line.strip()
        line = line.lstrip('v ')
        spt = line.split(' ')
        if 'nan' in line:
            continue
        if len(spt) == 6:
            pts.append(torch.tensor([float(x) for x in spt]))
        if len(spt) == 3:
            t = [float(x) for x in spt]
            if append_normals:
                t += [0.0 for _ in range(3)]
            pts.append(torch.tensor(t))

    rtn = torch.stack(pts, dim=0)
    return rtn


def read_xyz_file(path: Path):
    with open(path, 'r') as file:
        return xyz2tensor(file.read(), append_normals=True)


def embed_color(img: torch.Tensor, color, box_size=70):
    shp = img.shape
    D2 = [shp[2] - box_size, shp[2]]
    D3 = [shp[3] - box_size, shp[3]]
    img = img.clone()
    img[:, :3, D2[0]:D2[1], D3[0]:D3[1]] = color[:, :, None, None]
    if img.shape[1] == 4:
        img[:, -1, D2[0]:D2[1], D3[0]:D3[1]] = 1
    return img


def get_n_params(model):
    pp=0
    for p in list(model.parameters()):
        nn=1
        for s in list(p.size()):
            nn = nn*s
        pp += nn
    return pp


def xyz2tensor(txt, append_normals=False):
    pts = []
    for line in txt.split('\n'):
        line = line.strip()
        line = line.lstrip('v ')
        spt = line.split(' ')
        if 'nan' in line:
            continue
        if len(spt) == 6:
            pts.append(torch.tensor([float(x) for x in spt]))
        if len(spt) == 3:
            t = [float(x) for x in spt]
            if append_normals:
                t += [0.0 for _ in range(3)]
            pts.append(torch.tensor(t))

    rtn = torch.stack(pts, dim=0)
    return rtn


def read_xyz_file(path: Path):
    with open(path, 'r') as file:
        return xyz2tensor(file.read(), append_normals=True)