#!/usr/bin/env python

import numpy as np
import math

import empty
import yaml_ut
import dbg


def pos_trans_arr(arr, pos):
	return arr - np.array( [ pos.x, pos.y, pos.z ] )

def pos_trans(xs, ys, zs, pos):
	arr = np.array( [ xs, ys, zs ] ).T
	arr = pos_trans_arr( arr, pos )
	return arr.T

def pa_trans(xs, ys, zs, pa):
	(xs, ys, zs) = pos_trans( xs, ys, zs, pa.p )

	s = math.sin( -pa.a )
	c = math.cos( -pa.a )
	(xs, ys) = ( c * xs - s * ys, s * xs + c * ys )

	return [ xs, ys, zs ]

def pa_trans_arr(arr, pa):
	(xs, ys, zs) = arr.T
	(xs, ys, zs) = pa_trans( xs, ys, zs, pa )
	return np.array( [ xs, ys, zs ] ).T


def get_d2(arr, p):
	(xs, ys, zs) = ( arr - p ).T
	d2_lst = xs * xs + ys * ys + zs * zs
	return d2_lst


def get_near_bools(xs, ys, zs, dist):
	d2 = xs * xs + ys * ys + zs * zs
	bools = d2 <= dist * dist
	return bools

def get_near_bools_arr(arr, dist):
	(xs, ys, zs) = arr.T
	return get_near_bools( xs, ys, zs, dist )

def get_near_arr(arr, dist):
	bools = get_near_bools_arr( arr, dist )
	return arr[ bools, : ]


def to_bytes_head(b):
	h = np.int32( len( b ) ).tobytes()
	return h + b

def from_bytes_head(b):
	(h, b) = ( b[ : 4 ], b[ 4 : ] )
	n = np.frombuffer( h, dtype=np.int32 )[ 0 ]
	(b, t) = ( b[ : n ], b[ n : ] )
	return ( b, t )

def to_bytes_dtype(lst, dtype):
	b = np.array( lst, dtype=dtype ).tobytes()
	return to_bytes_head( b )

def from_bytes_dtype(b, dtype):
	(b, t) = from_bytes_head( b )
	arr = np.frombuffer( b, dtype=dtype ).copy()
	return (arr, t)

def to_bytes_yaml( o ):
	return to_bytes_head( yaml_ut.dump( o ).encode( 'utf-8' ) )

def from_bytes_yaml( b ):
	( b, t ) = from_bytes_head( b )
	o = yaml_ut.load( b.decode( 'utf-8' ) )
	return ( o, t )

def to_bytes_arr( arr ):
	dic = {}
	dic[ 'dtype' ] = str( arr.dtype )
	dic[ 'shape' ] = list( arr.shape )
	b = to_bytes_yaml( dic )
	b += to_bytes_head( arr.tobytes() )
	return b

def from_bytes_arr( b ):
	( dic, b ) = from_bytes_yaml( b )
	dtype = dic.get( 'dtype', 'float32' )
	shape = dic.get( 'shape' )
	( arr, t ) = from_bytes_dtype( b, dtype )
	if shape:
		shape = [ -1 ] + list( shape[ 1 : ] )
		arr = arr.reshape( shape )
	return ( arr, t )		

def save_bytes(path, b, show=True):
	end = dbg.start_save( path, show )
	with open( path, 'wb' ) as f:
		f.write( b )
	end()

def load_bytes(path, show=True):
	end = dbg.start_load( path, show )
	b = b''
	with open( path, 'rb' ) as f:
		b = f.read()
	end()
	return b

def save_bytes_dtype(path, lst, dtype, show=True):
	save_bytes( path, to_bytes_dtype( lst, dtype ), show )

def load_bytes_dtype(path, dtype, show=True):
	(arr, t) = from_bytes_dtype( load_bytes( path, show ), dtype )
	return (arr, t )

def write_bytes_head( f, b ):
	f.write( to_bytes_head( b ) )

def read_bytes_head( f ):
	h = f.read( 4 )
	if not h:
		return h
	n = np.frombuffer( h, dtype=np.int32 )[ 0 ]
	return f.read( n )


def get_inf(dic):
	name = dic.get( 'name' )
	is_list = 'type_list' in dic
	typ = dic.get( 'type_list' if is_list else 'type' )
	n = dic.get( 'n', 1 )
	nest = type( typ ) == list
	dtype = eval( typ ) if not nest else None
	dtype_sz = np.dtype( dtype ).itemsize if dtype else None

	def set_n(n):
		dic[ 'n' ] = n

	return empty.new( locals() )


def tobytes_dic(o, dic):
	inf = get_inf( dic )

	v = getattr( o, inf.name, None )
	arr = v if inf.is_list else None
	if arr:
		inf.set_n( len( arr ) )

	if inf.nest:
		dics = inf.typ
		f = lambda o: tobytes( o, dics )
		if not inf.is_list:
			return f( v )
		return b''.join( map( f, arr ) )

	if type( arr ) == bytes:
		return arr

	if not inf.is_list:
		arr = [ v ]

	if arr and type( arr ) != np.ndarray or type( arr[ 0 ] ) != inf.dtype:
		arr = np.array( arr, dtype=inf.dtype )
		inf.set_n( len( arr ) )

	return arr.tobytes()

def tobytes(o, dics):
	f = lambda dic: tobytes_dic( o, dic )
	return b''.join( map( f, dics ) )


def frombytes_dic(buf, dic):
	inf = get_inf( dic )

	if inf.nest:
		dics = inf.typ
		lst = []
		for i in range( inf.n ):
			(o, buf) = frombytes( buf, dics )
			lst.append( o )

		v = lst if inf.is_list else lst[ 0 ]
		return ( inf.name, v, buf )
			
	v = np.frombuffer( buf, dtype=inf.dtype, count=inf.n )
	n = inf.dtype_sz * inf.n
	buf = buf[ n: ]

	if not inf.is_list and inf.n == 1:
		v = v[ 0 ]

	if inf.dtype == np.uint8:
		v = v.tobytes()
	return ( inf.name, v, buf )

def frombytes(buf, dics):
	o = empty.new()

	for dic in dics:
		(name, v, buf) = frombytes_dic( buf, dic )
		setattr( o, name, v )

	return ( o, buf )


def save(path, o, s_yaml):
	dics = yaml_ut.load( s_yaml )

	buf = tobytes( o, dics )

	s_yaml = yaml_ut.dump_no_flow( dics ).decode()

	n = s_yaml.count( '\n' )
	s = str( n ) + '\n' + s_yaml

	with open( path, 'wb' ) as f:
		f.write( s.encode() )
		f.write( buf )

def info(path):
	buf = b''
	with open( path, 'rb' ) as f:
		s = f.readline()
		n = int( s )
		for i in range( n ):
			buf += f.readline()
		pos = f.tell()

	s_yaml = buf.decode()
	return ( s_yaml, pos )

def load(path):
	(s_yaml, pos) = info( path )

	buf = b''
	with open( path, 'rb' ) as f:
		f.seek( pos )
		buf = f.read()

	dics = yaml_ut.load( s_yaml )

	(o, buf) = frombytes( buf, dics )
	return o


# vec

arr_new = lambda *args: np.array( args )

p000 = arr_new( 0, 0, 0 )
v100 = arr_new( 1, 0, 0 )
v010 = arr_new( 0, 1, 0 )
v001 = arr_new( 0, 0, 1 )

v_len = lambda v: np.dot( v, v ) ** 0.5

def v_unit(v):
	d = v_len( v )
	return v / d if d != 0 else v001

cross_unit = lambda va, vb: v_unit( np.cross( va, vb ) )

cross_y = lambda vx: cross_unit( v001, vx )

cross_z = lambda vx: cross_unit( vx, cross_y( vx ) )


def get_dist_lst(ps, p):
	return ( ( ps - p ).T ** 2 ).sum( axis=0 ) ** 0.5


# q4

def q4_rev(q4):
	(x, y, z, w) = q4
	return arr_new( -x, -y, -z, w )

def q4_mat(q4):
	(x, y, z, w) = q4
	return np.array( [
		[ w, -z, y, x ],
		[ z, w, -x, y ],
		[ -y, x, w, z ],
		[ -x, -y, -z, w ],
	] )

def q4_rot(q4, q4_targ):
	return q4_mat( q4 ).dot( q4_targ )

def q4_rot_v(q4, v):
	v4 = arr_new( v[ 0 ], v[ 1 ], v[ 2 ], 0 )
	q4_ = q4_rev( q4 )
	r4 = q4_rot( q4_rot( q4, v4 ), q4_ )
	return r4[ : 3 ]  # [ x, y, z ]


def q4_cnv(targ_q4, base_q4):
	return q4_rot( q4_rev( base_q4 ), targ_q4 )

def pos_cnv(targ_pos, base_pos, base_q4):
	return q4_rot_v( q4_rev( base_q4 ), targ_pos - base_pos )

def pose_cnv(targ_pos, targ_q4, base_pos, base_q4):
	pos = pos_cnv( targ_pos, base_pos, base_q4 )
	q4 = q4_cnv( targ_q4, base_q4 )
	return ( pos, q4 )


def q4_icnv(targ_q4, base_q4):
	return q4_rot( base_q4, targ_q4 )

def pos_icnv(targ_pos, base_pos, base_q4):
	return q4_rot_v( base_q4, targ_pos ) + base_pos

def pose_icnv(targ_pos_on_base, targ_q4_on_base, base_pos, base_q4):
	pos = pos_icnv( targ_pos_on_base, base_pos, base_q4 )
	q4 = q4_icnv( targ_q4_on_base, base_q4 )
	return ( pos, q4 )


def q4_new_v_arg(v, arg):
	L = v_len( v )
	if L == 0:
		return arr_new( 0, 0, 0, 1 )
	s = math.sin( arg * 0.5 )
	(x, y, z) = v * ( s / L )
	w = math.cos( arg * 0.5 )
	return arr_new( x, y, z, w )

def q4_to_v_arg( q4 ):
	( x, y, z, w ) = q4
	v = arr_new( x, y, z )
	s = v_len( v )
	if s == 0:
		return ( v001, 0 )  # !
	c = w
	arg_h = math.atan2( s, c )
	return ( v_unit( v ), arg_h * 2 )

def get_arg( va, vb, v_rot=None ):
	if v_rot is not None:
		va = np.cross( v_rot, va )
		vb = np.cross( v_rot, vb )
	va = v_unit( va )
	vb = v_unit( vb )
	crs = np.cross( va, vb )	
	arg = math.atan2( v_len( crs ), np.dot( va, vb ) )
	if v_rot is not None and np.dot( crs, v_rot ) < 0:
		arg = -arg
	return arg

def std_arg( arg ):
	return math.atan2( math.sin( arg ), math.cos( arg ) )

def posi_arg( arg ):
	arg = std_arg( arg )
	if arg < 0:
		if abs( arg ) < 0.5 * math.pi / 180:
			return 0  # !
		arg += 2 * math.pi
	return arg

def rot_v( v_rot, arg, v_targ ):
	q4 = q4_new_v_arg( v_rot, arg )
	return q4_rot_v( q4, v_targ )

def rot_li_pos( li_rot, arg, p_targ ):
	( c, v_rot ) = li_rot
	return c + rot_v( v_rot, arg, p_targ - c )


def ypr_to_q4(yaw, pitch, roll):
	# yaw, pitch, roll: radian

	yaw_ = q4_new_v_arg( v001, yaw )
	pitch_ = q4_new_v_arg( v010, pitch )
	roll_ = q4_new_v_arg( v100, roll )

	return q4_rot( roll_, q4_rot( pitch_, yaw_ ) )

def q4_to_ypr(q4):
	v = q4_rot_v( q4, v001 )

	roll = 0
	(x, y, z) = v
	if y != 0 or z != 0:
		roll = math.atan2( -y, z )

	v = q4_rot_v( q4_new_v_arg( v100, -roll ), v )
	
	pitch = 0
	(x, y, z) = v
	if z != 0 or x != 0:
		pitch = math.atan2( x, z )

	v = v100
	v = q4_rot_v( q4, v )
	v = q4_rot_v( q4_new_v_arg( v100, -roll ), v )
	v = q4_rot_v( q4_new_v_arg( v010, -pitch ), v )

	yaw = 0
	(x, y, z) = v
	if x != 0 or y != 0:
		yaw = math.atan2( y, x )

	return ( yaw, pitch, roll )  # radian



def line_plane_cross(line_p, line_v, plane_p, plane_v):
	# ( ( line_p + L * line_v ) - plane_p ) dot plane_v == 0
	# ( line_p - plane_p ) dot plane_v + L * ( line_v dot plane_v ) == 0
	# L = ( plane_p - line_p ) dot plane_v / ( line_v dot plane_v )

	L = np.dot( plane_p - line_p, plane_v ) / np.dot( line_v, plane_v )
	return line_p + L * line_v

def line_plane_cross_point( li, pl ):
	( li_p, li_v ) = li
	( pl_p, pl_v ) = pl
	if np.dot( li_v, pl_v ) == 0:
		return None
	return line_plane_cross( li_p, li_v, pl_p, pl_v )

def plane2_cross_line( pl_a, pl_b ):
	( ap, av ) = pl_a
	( bp, bv ) = pl_b
	if all( av == bv ):
		return None
	v = cross_unit( av, bv )
	tv = cross_unit( av, v )
	p = line_plane_cross_point( [ ap, tv ], pl_b )
	if p is None:
		return None
	return [ p, v ]

def plane3_cross_point( pl3 ):
	li = plane2_cross_line( pl3[ 0 ], pl3[ 1 ] )
	if li is None:
		return None
	return line_plane_cross_point( li, pl3[ 2 ] )

def point3_circle_point( p3 ):
	( p0, p1, p2 ) = p3
	pl0 = [ ( p0 + p1 ) * 0.5, p1 - p0 ]
	pl1 = [ ( p1 + p2 ) * 0.5, p2 - p1 ]
	pl2 = [ p1, cross_unit( p0 - p1, p2 - p1 ) ]
	return plane3_cross_point( [ pl0, pl1, pl2 ] )

def point3_circle_touch_v( p3 ):
	c = point3_circle_point( p3 )
	if c is None:
		return None
	v3 = p3 - c
	rot_v = np.cross( v3[ 0 ], v3[ 1 ] )
	return list( map( lambda v: cross_unit( rot_v, v ), v3 ) )

def point_plane_len(p, plane_p, plane_v):
	return np.abs( np.dot( plane_p - p, plane_v ) / np.dot( plane_v, plane_v ) )

def get_plane_z( pl_p, pl_v, x, y ):
	return ( np.dot( pl_v, pl_p ) - pl_v[ 0 ] * x - pl_v[ 1 ] * y ) / pl_v[ 2 ]


def get_arg_from_line_p2(line_p, line_v, ps, pe):
	os = line_plane_cross( line_p, line_v, ps, line_v )
	oe = line_plane_cross( line_p, line_v, pe, line_v )

	vs = v_unit( ps - os )
	ve = v_unit( pe - oe )
	c = ( vs + ve ) * 0.5

	dx = v_len( c )
	dy = v_len( vs - c )
	arg = 2 * math.atan2( dy, dx )

	if np.dot( np.cross( vs, c ), line_v ) < 0:
		arg = -arg
	return arg

def get_q4_from_vx_vy(vx, vy):
	vx = v_unit( vx )
	vy = v_unit( vy )
	vz = cross_unit( vx, vy )

	dx = vx - v100
	dy = vy - v010
	dz = vz - v001

	lx = v_len( dx )
	ly = v_len( dy )
	lz = v_len( dz )

	if lx >= lz and ly >= lz:
		rot_v = cross_unit( dx, dy )
		arg_x = get_arg_from_line_p2( p000, rot_v, v100, vx )
		arg_y = get_arg_from_line_p2( p000, rot_v, v010, vy )
		arg = ( arg_x + arg_y ) * 0.5
	elif ly >= lx and lz >= lx:
		rot_v = cross_unit( dy, dz )
		arg_y = get_arg_from_line_p2( p000, rot_v, v010, vy )
		arg_z = get_arg_from_line_p2( p000, rot_v, v001, vz )
		arg = ( arg_y + arg_z ) * 0.5
	else:
		rot_v = cross_unit( dz, dx )
		arg_z = get_arg_from_line_p2( p000, rot_v, v001, vz )
		arg_x = get_arg_from_line_p2( p000, rot_v, v100, vx )
		arg = ( arg_z + arg_x ) * 0.5

	return q4_new_v_arg( rot_v, arg )

def corr_q4_by_ps_pe(ps, pe, targ_q4):
	vx = pe - ps
	inc_vz = pos_icnv( v001, p000, targ_q4 )
	vy = cross_unit( inc_vz, vx )
	base_q4 = get_q4_from_vx_vy( vx, vy )
	return q4_cnv( targ_q4, base_q4 )


def newton( f, a, b, lmt ):
	va = f( a )
	vb = f( b )
	if va * vb > 0:
		return None

	while abs( b - a ) > lmt:
		if va == 0:
			return a
		if vb == 0:
			return b

		c = ( a + b ) * 0.5
		vc = f( c )
		if va * vc >= 0:
			( a, va ) = ( c, vc )
		else:
			( b, vb ) = ( c, vc )
	return c


def curve_new(plst, vs=None, ve=None):
	# plst = [ [ x, y, z ], [ x, y, z], ... ]

	n = len( plst )
	Ls = np.zeros( n )
	Ds = np.zeros_like( plst )
	dlst = np.zeros_like( plst )

	for i in range( n - 1 ):
		ps = plst[ i ]
		pe = plst[ i + 1 ]
		Ls[ i ] = v_len( pe - ps)
		Ds[ i ] = ( pe - ps ) / Ls[ i ]

	for i in range( n - 2 ):
		ds = Ds[ i ]
		de = Ds[ i + 1 ]
		dlst[ i + 1 ] = ( ds + de ) * 0.5

	dlst[ 0 ] = Ds[ 0 ] * 2 - dlst[ 1 ] if vs is None else v_unit( vs )
	dlst[ -1 ] = Ds[ n - 2 ] * 2 - dlst[ -2 ] if ve is None else v_unit( ve )


	def get_ABCDL(ps, pe, ds, de, L):
		# 3A * LL + 2B L + C = de
		# A * LLL + B LL + C L + D = pe

		# 3A * LLL + 2B LL + C L = de L
		# 2A * LLL + 2B LL + 2C L + 2D = 2 pe

		# A * LLL - C L - 2 D = de L - 2 pe
		# A = ( de L - 2 pe + C L + 2 D ) / LLL
		# A = ( ( de + C ) * L + 2 * ( D - pe ) ) / LLL

		# 2B L = de - C - 3A * LL
		# B = ( de - C - 3A * LL ) / 2L

		C = ds
		D = ps
		A = ( ( de + C ) * L + 2 * ( D - pe ) ) / L ** 3
		B = ( de - C - 3 * A * L * L ) / ( 2 * L )
		return [ A, B, C, D, L ]

	def get_ABCDL_i( i ):
		return get_ABCDL( plst[ i ], plst[ i + 1 ], dlst[ i ], dlst[ i + 1 ], Ls[ i ] )

	ABCDLs = []
	for i in range( n - 1 ):
		ABCDL = get_ABCDL_i( i )
		ABCDLs.append( ABCDL )

	def update_ABCDLs( i ):
		if 0 <= i and i < len( ABCDLs ):
			ABCDLs[ i ] = get_ABCDL_i( i )

	def update_dlst( i, v ):
		# i < len( plst )
		dlst[ i ] = v
		update_ABCDLs( i )
		update_ABCDLs( i - 1 )

	def get_pos_lst( dist=0.5, strict_spc=False ):
		get_pos = lambda A, B, C, D, t: A * t ** 3 + B * t ** 2 + C * t + D

		def get_pos_idx_rate( i, rate ):
			( A, B, C, D, L ) = ABCDLs[ i ]
			t = L * rate
			return get_pos( A, B, C, D, t )

		lst = []
		n = len( ABCDLs )

		if strict_spc:
			( i, rate ) = ( 0, 0 )
			p = get_pos_idx_rate( i, rate )
			lst.append( p )
			while i < n:
				f = lambda r: v_len( get_pos_idx_rate( i, r ) - lst[ -1 ] ) - dist
				rate = newton( f, rate, 1.0, 0.001 )
				if rate is None:
					i += 1
					rate = 0
				else:
					p = get_pos_idx_rate( i, rate )
					lst.append( p )
			p = get_pos_idx_rate( n - 1, 1.0 )
			lst.append( p )
		else:
			for i in range( n ):
				(A, B, C, D, L) = ABCDLs[ i ]
				m = max( int( L / dist ), 1 )
				for j in range( m ):
					t = L * j / m
					lst.append( get_pos( A, B, C, D, t ) )
				if i == n - 1:
					lst.append( get_pos( A, B, C, D, L ) )
		return lst

	return empty.new( locals() )


def get_plane_from_points(ps):
	pl_p = np.mean( ps, axis=0 )

	adj = ps - pl_p
	mat = np.cov( adj.T )
	(w, v) = np.linalg.eig( mat )

	sort = w.argsort()[::-1]
	#w = w[ sort ]
	v = v[ : , sort ]

	pl_v = v[ : , 2 ]

	if pl_v.dtype == np.complex128:
		return None  # !

	return ( pl_p, pl_v )

# EOF