import argparse
import json
from math import *
import os
import sys

import numpy as np
from plyfile import PlyData
from tqdm import tqdm, trange


EPS = sys.float_info.epsilon


class Triangle:

    def __init__(self, v1, v2, v3):
        # Vertices
        self.v1 = v1
        self.v2 = v2
        self.v3 = v3
        # Sort triangle vertices so that so that the normal points "outward"
        if np.dot(self.v1, np.cross(self.v2 - self.v1, self.v3 - self.v1)) < 0:
            self.v3, self.v2 = self.v2, self.v3
        # Edges of the origin vertex and their L2 norms
        self.e1 = self.v2 - self.v1
        self.e2 = self.v3 - self.v1
        self.e1n = np.linalg.norm(self.e1)
        self.e2n = np.linalg.norm(self.e2)
        # Normal vector of the triangle's plane
        self.n = np.cross(self.e1 / self.e1n, self.e2 / self.e2n)
        self.n = self.n / np.linalg.norm(self.n)
        # The vertex coordinates in plane coordinates
        self.v1p = self.plane_coords(self.v1)
        self.v2p = self.plane_coords(self.v2)
        self.v3p = self.plane_coords(self.v3)

    def intersect(self, ray):
        '''
        :param ray: numpy.array(3) Ray cast from [0, 0, 0]
        :return: A tuple of [bool, Optional[numpy.array(3)]]:
                 The first component is True if the ray intersects the triangle, False otherwise.
                 The second component marks the intersection of the triangle's plane (even if outside the triangle), or
                 None if the ray does not intersect the plane at all.
        '''
        # https://en.wikipedia.org/wiki/Möller–Trumbore_intersection_algorithm
        ray_cross_e2 = np.cross(ray, self.e2)
        dot = np.dot(self.e1, ray_cross_e2)
        if abs(dot) < EPS:
            return False, None
        s = -self.v1
        u = np.dot(s, ray_cross_e2) / dot
        s_cross_e1 = np.cross(s, self.e1)
        v = np.dot(ray, s_cross_e1) / dot
        t = np.dot(self.e2, s_cross_e1) / dot
        if (u < 0 and abs(u) > EPS) or (u > 1 and abs(u - 1) > EPS):
            return False, ray * t
        if (v < 0 and abs(v) > EPS) or (u + v > 1 and abs(u + v - 1) > EPS):
            return False, ray * t
        if t <= EPS:
            return False, None
        return True, ray * t

    def plane_coords(self, v):
        u = np.dot(self.e1 / self.e1n, v - self.v1)
        n = np.cross(self.n, self.e1 / self.e1n)
        v = -np.dot(n, v - self.v1)
        return np.array([u, v])


class Face:

    def __init__(self, vs):
        self.vertices = vs
        self.tris = []
        avg = np.average(self.vertices, axis=0)
        self.center = convert_to_spherical(*avg)
        # Sort vertices so that so that the normal points "outward"
        if np.dot(self.vertices[0], np.cross(self.vertices[1] - self.vertices[0], self.vertices[2] - self.vertices[0])) < 0:
            self.vertices = self.vertices[::-1]
        for i in range(1, len(self.vertices)-1):
            self.tris.append(Triangle(self.vertices[0], self.vertices[i], self.vertices[i+1]))
        self.plane_coords = [self.tris[0].plane_coords(v) for v in self.vertices]


def convert_to_spherical(x, y, z):
    xy = sqrt(x**2+y**2)
    lon = atan2(z, xy)*180/pi
    lat = atan2(y, x)*180/pi
    mag = sqrt(x**2+y**2+z**2)
    return np.array([lon, lat, mag])


def convert_to_cartesian(lon, lat):
    x = cos(lon*pi/180)
    y = sin(lon*pi/180)
    z = tan(lat*pi/180)
    mag = sqrt(x**2 + y**2 + z**2)
    return np.array([x/mag, y/mag, z/mag])


def map_poly(tri, poly, ref=None):
    if ref is None:
        ref = tri
    p0 = convert_to_cartesian(*poly[0])
    p0b, p0i = tri.intersect(p0)
    mapped = []
    new = True
    for i in range(1, len(poly)):
        p1 = convert_to_cartesian(*poly[i])
        p1b, p1i = tri.intersect(p1)
        if p0b or p1b and p0i is not None and p1i is not None:
            if new:
                mapped.append([ref.plane_coords(p0i), ref.plane_coords(p1i)])
                new = False
            else:
                mapped[-1].append(ref.plane_coords(p1i))
        else:
            new = True
        p0 = p1
        p0i = p1i
        p0b = p1b
    return mapped


def main(ns):
    with open(ns.faces, 'rb') as f:
        plydata = PlyData.read(f)
    vs = plydata['vertex']
    fs = plydata['face']['vertex_indices']
    faces = [Face([np.array(vs[vi].tolist())*ns.scale for vi in fvs]) for fvs in fs]
    with open(ns.geojson, 'r') as f:
        geojson = json.load(f)

    for i, face in tqdm(list(enumerate(faces)), desc='Faces   '):
        countries = {}
        for t in face.tris:
            for f in tqdm(geojson['features'], desc='Features', total=len(face.tris)*len(geojson['features'])):
                cc = f['properties']['ADMIN']
                if f['geometry']['type'] == 'MultiPolygon':
                    for poly in f['geometry']['coordinates']:
                        countries.setdefault(cc, []).extend(map_poly(t, poly[0], face.tris[0]))
                elif f['geometry']['type'] == 'Polygon':
                    countries.setdefault(cc, []).extend(map_poly(t, f['geometry']['coordinates'][0], face.tris[0]))
        minx = min(v[0] for v in face.plane_coords)
        miny = min(v[1] for v in face.plane_coords)
        maxx = max(v[0] for v in face.plane_coords)
        maxy = max(v[1] for v in face.plane_coords)
        w = maxx - minx
        h = maxy - miny
        formatters = dict(
            faces=os.path.basename(ns.faces).rsplit('.', 1)[0],
            geojson=ns.geojson,
            face=i,
            lon=face.center[0],
            lat=face.center[1],
            mag=face.center[2],
            ilon=int(face.center[0]),
            ilat=int(face.center[1]),
            imag=int(face.center[2]),
        )
        outfile = ns.out.format(**formatters)
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        with open(outfile, 'w') as svg:
            svg.write(f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="{minx} {miny} {w} {h}" width="{w}" height="{h}" id="face{i}">\n')
            svg.write(f'  <g id="cut">\n    <path id="outline" fill="none" stroke="black" stroke-width="{0.01*ns.scale}" d="M {face.plane_coords[0][0]} {face.plane_coords[0][1]}')
            for v in face.plane_coords[1:]:
                svg.write(f' L {v[0]} {v[1]}')
            svg.write(' Z" />\n  </g>\n  <g id="engrave">\n')
            for c, l in countries.items():
                if l:
                    svg.write(f'    <g id="{c}">\n')
                    for ls in l:
                        svg.write(f'      <path fill="none" stroke="red" stroke-width="{0.01*ns.scale}" d="M {ls[0][0]} {ls[0][1]}')
                        for x, y in ls[1:]:
                            svg.write(f' L {x} {y}')
                        svg.write('" />\n')
                    svg.write('    </g>\n')
            svg.write('  </g>\n</svg>\n')


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--geojson', '-g', default='countries.geojson')
    ap.add_argument('--faces', '-f', default='shapes/cube.json')
    ap.add_argument('--out', '-o', default='faces/{faces}/face{face}_{ilat}_{ilon}.svg')
    ap.add_argument('--scale', '-s', type=float, default=1)
    ns = ap.parse_args()
    main(ns)