# This file is part of T-Mart.
#
# Copyright 2023 Yulun Wu.
#
# T-Mart is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

import pandas as pd
import numpy as np
import sys
from .tm_geometry import angle_3d, linear_distance
from copy import copy

# intersect photon with atmosphere 

# new find_atm two scattering OTs for travelling through multiple layers 
def find_atm2(atm_profile,q1):
    
    z = q1[2]/1000
    
    if z < 0:
        print("ERROR: z < 0")
        return None
    
    alt_top_minus_z = atm_profile[:,1] - z
    
    # index of minimum positive difference 
    idx = np.where(alt_top_minus_z >= 0, alt_top_minus_z, np.inf).argmin()
    
    ot_rayleigh = atm_profile[idx,3]
    ot_mie = atm_profile[idx,4]
    
    return ot_rayleigh, ot_mie
    
# Testing intersecting triangles 
# 2: test boxes first 
# This one is a lot faster 
def intersect_line_DEMtri2(q0, q1, DEM_tri, print_on = False):
    '''

    Parameters
    ----------
    q0 : list
        starting point of the line.
    q1 : list
        ending point.
    DEM_tri : list
        a list of numpy arrays.

    Returns
    -------
    intersect_tri : pandas dataframe
        collision coordinates, direction normal to the surface
        and linear distance to the starting point.

    '''
    
    # Manual switch 
    print_on = False 
    
    ### Identify boxes intersecting the line first     

    tri_x = DEM_tri[0][:,0] 
    tri_y = DEM_tri[0][:,1] 
    tri_z = DEM_tri[0][:,2]
    
    tri_x_min = np.min(tri_x,axis=0)
    tri_x_max = np.max(tri_x,axis=0)

    tri_y_min = np.min(tri_y,axis=0)  
    tri_y_max = np.max(tri_y,axis=0)    
    
    # Find the box within max z and min z, project it onto XY and select the pixels 
    pt_direction_c = [q1[0]-q0[0],q1[1]-q0[1],q1[2]-q0[2]]
    z_max = np.max(tri_z) + 0.1 
    z_min = np.min(tri_z) - 0.1
    pt_direction_c = q1 - q0
    
    # make box 
    q_z_max = q0 + ( (z_max - q0[2]) / pt_direction_c[2] ) * pt_direction_c
    q_z_min = q0 + ( (z_min - q0[2]) / pt_direction_c[2] ) * pt_direction_c
    
    # First treat all triangles as boxes (max/min XYZ)!!!
    box_x_min = min(q_z_max[0],q_z_min[0])
    box_x_max = max(q_z_max[0],q_z_min[0])
    box_y_min = min(q_z_max[1],q_z_min[1])
    box_y_max = max(q_z_max[1],q_z_min[1])
    
    # boxes that cross z
    crossing_z = np.logical_and(np.logical_and(tri_x_min < box_x_max, tri_x_max > box_x_min),
                                np.logical_and(tri_y_min < box_y_max, tri_y_max > box_y_min))
    
    # crossing is for all xyz 
    if np.sum(crossing_z) <= 1:
        crossing = crossing_z
        
    else:
        if q0[0] == q1[0] and q0[1] == q1[1]:
            if print_on: print('Vertical move')
            crossing = crossing_z
            
        elif q0[0] == q1[0]:
            if print_on: print('Equal x, move along y')
            crossing_y = np.logical_and(tri_x_min < q0[0], tri_x_max > q0[0])
            crossing = np.logical_and(crossing_z, crossing_y)
            
        elif q0[1] == q1[1]:
            if print_on: print('Equal y, move along x')    
            crossing_x = np.logical_and(tri_y_min < q0[1], tri_y_max > q0[1])
            crossing = np.logical_and(crossing_z, crossing_x)
            
        else:
            if print_on: print('Move along both x and y')
            
            # test if intersects the top or bottom of the square
            
            # equation on the x-y plane
            slope = (q1[1] - q0[1]) / (q1[0] - q0[0])
            y_itcp = q1[1] - q1[0]*slope
            
            # crossing the upper boundary (more negative y)
            x_itcp_y_min = (tri_y_min - y_itcp) / slope
            
            # crossing the lower boundary (more positive y)
            x_itcp_y_max = (tri_y_max - y_itcp) / slope
            
            y_itcp_x_min = tri_x_min * slope + y_itcp
            y_itcp_x_max = tri_x_max * slope + y_itcp

            crossing_upper = np.logical_and(tri_x_min < x_itcp_y_min, x_itcp_y_min < tri_x_max)
            crossing_lower = np.logical_and(tri_x_min < x_itcp_y_max, x_itcp_y_max < tri_x_max)
            crossing_x = np.logical_or(crossing_upper,crossing_lower) 
            
            
            crossing_left  = np.logical_and(tri_y_min < y_itcp_x_min, y_itcp_x_min < tri_y_max)
            crossing_right = np.logical_and(tri_y_min < y_itcp_x_max, y_itcp_x_min < tri_y_max)
            crossing_y = np.logical_or(crossing_left,crossing_right) 
            
            
            crossing_xy = np.logical_or(crossing_x,crossing_y) 
            crossing = np.logical_and(crossing_z, crossing_xy)

    # intersecting triangles 
    intersect_tri = pd.DataFrame() 
    
    for tri in DEM_tri: # tri is a set of triangle, ref_tri1[point0-2, xyz:0-2, row, column]
    
        for i in range(0,tri[:,:,crossing].shape[2]):

            p0 = tri[:,:,crossing] [0,:,i]
            p1 = tri[:,:,crossing] [1,:,i]
            p2 = tri[:,:,crossing] [2,:,i]
                
            # Extract coordinates of intersection if exists 
            intersect = _intersect_line_triangle(q0,q1,p0,p1,p2)

            
            if (intersect is not None):             
                
                if print_on:
                    print("Line intersecting triangle at: " + str(intersect))
                
                N = np.cross(p1-p0, p2-p0) 
                
                # test normal if in the same direction as the incoming line 
                # if not, reverse normal direction 
                
                if angle_3d(q0,intersect,(intersect + N)) > 90: N = -N
                
                intersect_tri_temp = pd.DataFrame({
                    # coordinates and direction 
                    'X': [intersect[0]],
                    'Y': [intersect[1]],
                    'Z': [intersect[2]],
                    'N_X': [N[0]],
                    'N_Y': [N[1]],
                    'N_Z': [N[2]]
                    })
                
                intersect_tri_temp['linear_distance'] = linear_distance(
                    intersect_tri_temp.iloc[0,0:3].tolist(), 
                    q0)
                
                # intersect_tri = intersect_tri.append(intersect_tri_temp) # append to concat on Feb 4, 2022
                intersect_tri = pd.concat([intersect_tri,intersect_tri_temp])
      
    intersect_tri.reset_index(drop=True, inplace=True)
    return intersect_tri

def _intersect_line_triangle(q1,q2,p1,p2,p3):
    '''

    Parameters
    ----------
    q1 : list
        Coordinates of one end of the line: X, Y, Z.
    q2 : list
        Coordinates of the other end of the line.
    p1 : list
        Coordinates of the first angle of the triangle.
    p2 : list
        Coordinates of the second angle of the triangle.
    p3 : list
        Coordinates of the third angle of the triangle.

    Returns
    -------
    Coordinates of the intersection if they intersect, or 
    None if they don't.


    '''
    # from https://stackoverflow.com/questions/42740765/intersection-between-line-and-triangle-in-3d
    def signed_tetra_volume(a,b,c,d):
        return np.sign(np.dot(np.cross(b-a,c-a),d-a)/6.0)

    s1 = signed_tetra_volume(q1,p1,p2,p3)
    s2 = signed_tetra_volume(q2,p1,p2,p3)

    if s1 != s2:
        s3 = signed_tetra_volume(q1,q2,p1,p2)
        s4 = signed_tetra_volume(q1,q2,p2,p3)
        s5 = signed_tetra_volume(q1,q2,p3,p1)
        if s3 == s4 and s4 == s5:
            n = np.cross(p2-p1,p3-p1)
            t = np.dot(p1-q1,n) / np.dot(q2-q1,n)
            return q1 + t * (q2-q1)
    return None

def reflectance_intersect(q_collision, image_reflectance, cell_size, bg_ref, bg_coords):
    '''
    Find the reflectance of a triangle when collision happens
    June 10, 2021: also if is water 

    Parameters
    ----------
    q_collision : TYPE
        DESCRIPTION.
    image_reflectance : TYPE
        DESCRIPTION.
    cell_size : TYPE
        DESCRIPTION.
    b1_ref : TYPE
        background reflectance 1.
    b2_ref : TYPE
        background reflectance 2.
    b1_a : TYPE
        first point dividing b1_ref and b2_ref.
    b1_b : TYPE
        2nd point.

    Returns
    -------
    triangle_reflectance : TYPE
        DESCRIPTION.

    '''
    
    b1_ref = bg_ref[0]
    b2_ref = bg_ref[1]
    
    image_shape = image_reflectance.shape

    b1_a, b1_b = np.array(bg_coords[0]), np.array(bg_coords[1])
    
    triangle_y = int(q_collision[1]//cell_size)
    triangle_x = int(q_collision[0]//cell_size)
         
    line = b1_a - b1_b
    slope = line[1]/line[0]  # y = slope * x + x_intercept 
    x_intercept = b1_a[1] - slope *b1_a[0]   
    
    # if on the padding triangles 
    if (triangle_y < 0 or triangle_x < 0 or 
        triangle_y >= image_shape[0] or triangle_x >= image_shape[1]):
        above_line = q_collision[1] >= (q_collision[0] * slope + x_intercept) 
        if x_intercept>=0:
            if above_line:
                triangle_reflectance = b2_ref
            else:
                triangle_reflectance = b1_ref
        else:
            if above_line:
                triangle_reflectance = b1_ref
            else:
                triangle_reflectance = b2_ref
    else: # if within the original image 
        triangle_reflectance =  image_reflectance[triangle_y, triangle_x]
        
    return triangle_reflectance

def reflectance_background(q_collision,bg_ref, bg_coords):
    # Find reflectance when outside the triangles 
    # June 10, 2021: also if is water 
    
    b1_ref = bg_ref[0]
    b2_ref = bg_ref[1]
    
    b1_a, b1_b = np.array(bg_coords[0]), np.array(bg_coords[1])
    line = b1_a - b1_b
    slope = line[1]/line[0]  # y = slope * x + x_intercept 
    x_intercept = b1_a[1] - slope *b1_a[0]  
    
    above_line = q_collision[1] >= (q_collision[0] * slope + x_intercept) 
        
    if x_intercept>=0:
        #print('below the line is b1')
        
        if above_line:
            triangle_reflectance = b2_ref
        else:
            triangle_reflectance = b1_ref
            
    else:
        #print('above the line is b1')
        
        if above_line:
            triangle_reflectance = b1_ref
        else:
            triangle_reflectance = b2_ref
    return triangle_reflectance


def intersect_background(q0,q1,bg_elevation):
    '''
    Find the XYZ of q0-q1 at Z 

    '''

    ratio = (q0[2]-bg_elevation) / (q0[2]-q1[2])
    x = q0[0] - ((q0[0]-q1[0])) * ratio
    y = q0[1] - ((q0[1]-q1[1])) * ratio

    return [x,y,bg_elevation]


# test if line intersects atm boundaries 
def intersect_line_boundary(q0,q1,atm_profile,print_on=False):
    '''
    return the closest layer 

    '''
    
    z0, z1 = q0[2]/1000,q1[2]/1000
    
    if z0==z1:
        sys.exit("intersect_line_boundary: z0==z1")
    
    going_up = z1>z0
    cross_b = False # if cross a boundary 
    out = False # if out
    
    # output 
    intersect_b = pd.DataFrame() # default output if no crossing boundary 
    
    if going_up:
        diff = atm_profile.Alt_bottom - z0
        if any(diff>0): # at least one alt_bottom above z0
            layer_idx = diff[diff>0].idxmin()
            
            # print(layer_idx)
            
            layer_bottom = atm_profile.Alt_bottom[layer_idx]
            # print(layer_bottom)
            
            if (z1 >= layer_bottom):
                # append info 
                cross_b = True 
                altitude = layer_bottom # altitude of crossing boundary 
        else: # all alt_bottom below z0 
            if (z1>=100 and going_up):
                out = True 
                cross_b = True
                altitude = 100
                layer_idx = atm_profile.Alt_bottom.idxmax() # return top layer because OTs are not used 
            
    else: # going down 
        diff = z0 - atm_profile.Alt_top
        
        if any(diff>0): # no alt_top above q0 --> no need to edit atm 
            layer_idx = diff[diff>0].idxmin()
            layer_top = atm_profile.Alt_top[layer_idx]
            if(z1 <= layer_top):
                # append info 
                cross_b = True 
                altitude = layer_top # altitude of crossing boundary 
            
    if cross_b: 
        intersect_b = copy(atm_profile.iloc[layer_idx])
        intersect_b['out'] = False
        
        # where crossing boundary  
        xyz = intersect_background(q0,q1,altitude*1000) + [altitude*1000]
        
        intersect_b['X'] = xyz[0]
        intersect_b['Y'] = xyz[1]
        intersect_b['Z'] = xyz[2]
        
        intersect_b['linear_distance'] = linear_distance(q0,xyz)

        if out:
            intersect_b['out'] = True
    
    return intersect_b 
