import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection


def pol2car(rho, pi, xi, yi):
    x = rho * np.cos(pi) + xi
    y = rho * np.sin(pi) + yi
    return (x, y)

def plot_isovist(isovists, show_axis=False, s=0.1, figsize=(5,5)):
    #transpose the matrix
    # isovists = np.transpose(isovists, (isovists.ndim-1, isovists.ndim-2))
    plt.switch_backend('agg')
    fig = plt.figure(figsize=figsize)
    points = []
    res = np.pi/90
    isovist = isovists
    for j, rho in enumerate(isovist):
        if rho < 1.0:
            pt = pol2car(rho, j*res, 0, 0)
            points.append(pt)
    x = [i[0] for i in points]
    y = [i[1] for i in points]
    ax = fig.add_subplot(111)
    ax.set_aspect('equal')
    ax.set_xlim(-1,1)
    ax.set_ylim(-1,1)
    if not show_axis:
        ax.axis('off')
    ax.scatter(x, y, s, 'black')
    return fig


def isovist_to_cartesian_a(isovist, x, y, scale):
    points = []
    xy = (x, y)
    res = np.pi/len(isovist)*2
    isovist = isovist * scale
    for j, rho in enumerate(isovist):
        pt = pol2car(rho, j*res, xy[0], xy[1])
        points.append(pt)
    points = np.stack(points)
    return(points)


def isovist_to_cartesian_segment_a(isovist, x, y, scale, max=0.98, min = 0.1, d=0.1):
    points = []
    segment = []
    xy = (x, y)
    res = np.pi/len(isovist)*2
    isovist = isovist * scale
    p_rho = isovist[-1]
    for j, rho in enumerate(isovist):
        delta = abs(p_rho-rho)
        if j == 0:
            first_rho = rho
        if rho < max * scale and rho > min * scale and delta < d * scale:
            pt = pol2car(rho, j*res, xy[0], xy[1])
            segment.append(pt)
        else:
            points.append(segment)
            segment = []
        p_rho = rho
    if first_rho < max * scale and  first_rho > min * scale and abs(rho-first_rho)< d * scale :
        if len(points) > 0:
            segment.extend(points[0])
            points[0]=segment
        else:
            points.append(segment)
    else:
        points.append(segment)
    segments = []
    for i in range(len(points)):
        if len(points[i])>0:
            segment = np.stack(points[i])
            segments.append(segment)
    return(segments)


def index_to_loc_grid(idx, d):
    if idx == 0:
        return np.array((0., 0.), dtype=np.float32)
    elif idx == 1:
        return np.array((d, 0.), dtype=np.float32)
    elif idx == 2:
        return np.array((d, d), dtype=np.float32)
    elif idx == 3:
        return np.array((0., d), dtype=np.float32)
    elif idx == 4:
        return np.array((-d, d), dtype=np.float32)
    elif idx == 5:
        return np.array((-d, 0.), dtype=np.float32)
    elif idx == 6:
        return np.array((-d, -d), dtype=np.float32)
    elif idx == 7:
        return np.array((0., -d), dtype=np.float32)
    elif idx == 8:
        return np.array((d, -d), dtype=np.float32)
    else:
        raise NameError('Direction unknown')



# showing isovist sequence grid
def seq_show_grid(locs, isovists, d=0.2, figsize=(8, 8), center=False, lim=1.5, alpha=0.02, rad=0.9, b_width=1.0, calculate_lim=False):
    # walk trough the sequence
    p_loc = np.array((0, 0))
    b_segments = []
    b_points = []
    isovists_pts = []
    res = np.pi/128
    cartesian_locs = []
    for loc, isovist in zip(locs, isovists):
        rel_pos = index_to_loc_grid(loc, d) + p_loc
        for j, rho in enumerate(isovist):
            if rho < rad :
                pt = pol2car(rho, j*res, rel_pos[0], rel_pos[1])
                b_points.append(pt)
        segments = isovist_to_cartesian_segment_a(isovist, rel_pos[0], rel_pos[1], 1.0)
        b_segments.extend(segments)
        isovists_pts.append(isovist_to_cartesian_a(isovist, rel_pos[0], rel_pos[1], 1.0))
        cartesian_locs.append(rel_pos)
        p_loc = rel_pos

    if len(b_points) > 0:
        b_points = np.stack(b_points)
    else:
        b_points =[]
    isovists_pts = np.stack(isovists_pts)
    # b_segments = np.stack(b_segments)
    cartesian_locs = np.stack(cartesian_locs)

    # set graphic properties
    isovist_path_width = 0.1
    isovist_path_pt1 = 6.0
    isovist_path_pt2 = 10.0
    isovist_boundary_pt = 0.05

    if center == True:

        bbox = get_bbox(b_points)
        center_pt = get_center_pts(bbox, np_array=True)
        b_points = [ pt - center_pt for pt in b_points]
        isovists_pts = [ pt - center_pt for pt in isovists_pts]
        b_segments =  [ pt - center_pt for pt in b_segments]
        cartesian_locs = [ pt - center_pt for pt in cartesian_locs]

    # resize image
    if calculate_lim == True:
        if bbox is not None:
            max = np.max(np.abs(bbox))
        else:
            max = 2.0
        if max > 2.0:
            lim = ((max // 0.5) + 1) * 0.5
            isovist_path_width *= 2.0/lim
            isovist_path_pt1 *= 2.0/lim
            isovist_path_pt2 *= 2.0/lim
            isovist_boundary_pt *= 2.0/lim


    fig, ax = plt.subplots(1,1, figsize=figsize, dpi=96)

    

    # isovists
    isovist_poly = []
    for isovist_pts in isovists_pts:
        isovist_poly.append(Polygon(isovist_pts, True))
    r = PatchCollection(isovist_poly)
    r.set_facecolor('#00aabb')
    r.set_edgecolor(None)
    r.set_alpha(alpha)
    ax.add_collection(r)

    

    # isovist path
    q = PatchCollection([Polygon(cartesian_locs, False)])
    q.set_facecolor('none')
    q.set_edgecolor('red')
    q.set_linewidth(isovist_path_width)
    # q.set_linestyle('dashed')
    ax.add_collection(q)
    
    # start_pt
    ax.scatter([x[0] for x in cartesian_locs[:1]], [x[1] for x in cartesian_locs[:1]], s = isovist_path_pt1, c='k', marker='s')

    # sequence
    ax.scatter([x[0] for x in cartesian_locs[1:-1]], [x[1] for x in cartesian_locs[1:-1]], s = isovist_path_pt1, c='red')

    # end pt
    ax.scatter([x[0] for x in cartesian_locs[-1:]], [x[1] for x in cartesian_locs[-1:]], s = isovist_path_pt2, c='k', marker='x')

    # boundaries
    edge_patches = []
    for segment in b_segments:
        if len(segment) > 5:
            polygon = Polygon(segment, False)
            edge_patches.append(polygon)
    p = PatchCollection(edge_patches)
    p.set_facecolor('none')
    p.set_edgecolor('#000000')
    p.set_linewidth(b_width)
    ax.scatter([x[0] for x in b_points], [x[1] for x in b_points], s = isovist_boundary_pt, c='#000000',)
    # ax.add_collection(p)


    # style
    ax.set_aspect('equal')
    lim = lim
    ax.set_xlim(-lim,lim)
    ax.set_ylim(-lim,lim)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.axis('off')

    return fig

# plotting isovist sequence grid
def plot_isovist_sequence_grid(locs, isovists, figsize=(8,8), center=False, lim=1.5, alpha=0.02, rad=0.9, b_width=1.0, calculate_lim=False):
    fig = seq_show_grid(locs, isovists, figsize=figsize, center=center, lim=lim, alpha=alpha, rad=rad, b_width=b_width, calculate_lim=calculate_lim)
    # for plot with torchvision util
    fig.canvas.draw()
    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    w, h = fig.canvas.get_width_height()
    im = data.reshape((int(h), int(w), -1))
    im = im.transpose((2, 0, 1))
    plt.close()
    return im

def get_bbox(pts):
    if len(pts) > 0:
        if type(pts) is list:
            pts = np.stack(pts)
        bbox = np.min(pts[:, 0]), np.max(pts[:, 0]), np.min(pts[:, 1]), np.max(pts[:, 1])
        return bbox
    else:
        return None

def get_center_pts(bbox, np_array = False):
    if bbox is not None:
        center = 0.5*(bbox[0] + bbox[1]),  0.5*(bbox[2] + bbox[3])
        if np_array:
            center = np.asarray(center)
    else:
        center = np.asarray([0,0])
    return center