import numpy as np
from scipy.optimize import curve_fit
import utils
import sys

data   = utils.load( sys.argv[1] )

"""
find fill points
"""
points = []

points += utils.find_point( data[ : , 0 ] , 0 ) # x_min
points += utils.find_point(
        data[ : , data.shape[1] - 1 ],
        data.shape[1] - 1            ,
) # x_max

index_min = 0
while data.shape[0] - 1 > index_min:
    index_min += 1
    if len( utils.find_point( data[ index_min , : ] , index_min , 'y' ) ) == 3:
        break

points.append(
    utils.find_point( data[ index_min , : ] , index_min , 'y' )[1]
) # y_min

index_max = data.shape[0] - 1
while index_min < index_max:
    index_max -= 1
    if len( utils.find_point( data[ index_max , : ] , index_max , 'y' ) ) == 3:
        break


points.append(
    utils.find_point( data[ index_max , : ] , index_max , 'y' )[1]
) # y_max

small_data = utils.compress( data , 5 )

points = utils.big_to_small( points , 5 )

# size - 1
points[ 1 ][ 1 ] -= 1
points[ 3 ][ 0 ] -= 1

# little shift to be inside the light
points[ 2 ][ 1 ] += 3
points[ 3 ][ 1 ] += 3

"""
fill data
"""

extremum = []
for point in points:
    taken_points = utils.small_to_big(
        utils.fill( small_data , point , 1000 ),
        5
    )
    extremum.append( [
        np.min( taken_points[ : , 1 ] ),
        np.max( taken_points[ : , 1 ] ),
        np.min( taken_points[ : , 0 ] ),
        np.max( taken_points[ : , 0 ] ),
    ] )

border = {
    'x': {
        'min': extremum[0][1] + 1,
        'max': extremum[1][0]    ,
    },
    'y': {
        'min': extremum[2][3] + 1,
        'max': extremum[3][2]    ,
    },
}

"""
label deletion
"""

mean_data = np.mean( data[
    border['y']['min'] : border['y']['max'],
    border['x']['min'] : border['x']['max']
] , axis = 0 )

gauss        = lambda x , sigma , mu , a , b : a * (
    1 / sigma * np.sqrt(
        2 * np.pi
    ) * np.exp(
        - ( x - mu ) ** 2 / ( 2 * sigma ** 2 )
    )
) + b
abciss       = np.arange(
        border['x']['min'],
        border['x']['max'],
        1
)
guess_params = [
    1                                                                    ,
    border['x']['min'] + ( border['x']['max'] - border['x']['min'] ) // 2,
    np.max( mean_data )                                                  ,
    np.min( mean_data )                                                  ,
]

first_estimate = curve_fit(
    gauss       ,
    abciss      ,
    mean_data   ,
    guess_params
)[0]

part_data   = [
    mean_data[ : mean_data.shape[0] // 2 ],
    mean_data[ mean_data.shape[0] // 2 : ]
]
part_abciss = [
    abciss[ : abciss.shape[0] // 2 ],
    abciss[ abciss.shape[0] // 2 : ]
]
part_result = []
for i in range( 2 ):
    part_result.append(
        curve_fit(
            gauss         ,
            part_abciss[i],
            part_data[i]  ,
            first_estimate
        )
    )

cov = np.array( [
    np.sum( np.diag( part_result[i][1] ) ) for i in range( 2 )
] )
i   = np.argmax( cov ) # part where the label is

derivee = np.convolve(
        np.gradient( part_data[i] ),
        np.ones( part_data[i].shape[0] // 100 ),
        'same',
)
start_label = np.argmax( derivee )
end_label   = np.argmin( derivee[ start_label :: ( - 1 ) ** i ] )

keys = [ 'min' , 'max' ]

border['x'][keys[i]] += ( - 1 ) ** i * ( start_label + end_label )

import matplotlib.pyplot as plt
plt.imshow( data[
    border[ 'y' ][ 'min' ] : border[ 'y' ][ 'max' ],
    border[ 'x' ][ 'min' ] : border[ 'x' ][ 'max' ]
] )
plt.savefig( 'asset/test.png' )