from src.core.afm import AFM as AFM
from src.core.afm import CHECKER as CHECKER
import src.core.analysis
import src.core.report
import src.core.tiff 
from time import sleep
import numpy as np
import os


class Wafer_Zero_Scan():
    
    def __init__(self):
        self.afm = AFM()
        self.base_dir = ''
        self.sub_dir = 'Wafer Zero Scan'
        self.file_name = 'RAW'
        self.parameter = None
    
    def move_to(self):
        pass


    def set(self):
        if self.parameter['xy_servo_enable'] == 'On':
            self.sub_dir += '/XYservo_On'
            self.parameter['z_scanner_range'] = 70
            self.parameter['xy_scanner_range'] = 100
            self.afm.enable_xy_servo()
       
        elif self.parameter['xy_servo_enable'] == 'Off':
            self.sub_dir += '/XYservo_Off'
            self.log_name = 'Position_XY.csv' 
            self.parameter['z_scanner_range'] = 20
            self.parameter['xy_scanner_range'] = 20
            self.afm.enable_xy_servo('off')
            self.header = 'Repeat, Lable , Input X Pos., Input Y Pos., Output X Pos., Output Y Pos.\n '

        if self.parameter['lift_enable'] == 'On':
            self.sub_dir += '_Lift_On'
        elif self.parameter['lift_enable'] == 'Off':
            self.sub_dir += '_Lift_Off'

        dir_path = self.base_dir + '\\' + self.sub_dir
        self.log_path = dir_path + '\\' + self.log_name
        if not os.path.exists(dir_path):os.makedirs(dir_path)
        fd = open(self.log_path,'w')
        fd.write(self.header)
        fd.close()


        self.afm.set_data_location(self.base_dir, self.sub_dir, self.file_name)
        self.afm.set_head_mode('contact')
        self.afm.clear_channels()
        self.afm.add_channel('ChannelZDriveOrTopography')
        
        self.afm.stop_scan()
        self.afm.set_z_scanner_range(self.parameter['z_scanner_range'])
        self.afm.set_xy_scanner_range(self.parameter['xy_scanner_range'])
        self.bck_scan_geometry = self.afm.get_scan_geometry_all()
        self.bck_scan_option = self.afm.get_scan_option_all()
        self.bck_approach_option = self.afm.get_approach_option_all()
        self.afm.set_scan_geometry(self.parameter['pixels'][0],self.parameter['pixels'][1], 0, 0)
        self.afm.set_scan_option(sine_scan=False,over_scan=False, \
                        over_scan_percent=0, tow_way=False, \
                        det_driven=False, force_slope_correction=False, \
                        interlace=False,slow_scan='off')
        self.afm.set_scan_rate(self.parameter['scan_rate'])
        self.afm.enable_z_servo()
        self.afm.set_z_servo(self.parameter['gain'], self.parameter['setpoint'])
        self.afm.set_approach_option()
    


    def unset(self):    
        self.afm.set_scan_option_dict(self.bck_scan_option)
        self.afm.set_scan_geometry_dict(self.bck_scan_geometry)
        self.afm.set_scan_option_dict(self.bck_scan_option)
    
    def __approach(self):
        self.afm.start_approach('q+s')
        wait_time = self.parameter['post_approach_wait']
        print(f"Waiting for {wait_time} seconds")
        sleep(wait_time)
    
    def __run(self):
        self.afm.trigger_image_scan()
        isscan = True
        while isscan:
            test = self.afm.query_scan_status()
            if test == 'true':isscan=True;
            else:isscan=False
            sleep(2)

    def approach(self):
        pass



    def run(self):
        pos = np.array(self.parameter['target_positions'])  # 'target_positions'에서 좌표 데이터 로드
        
        pos_int = []   # 문자열 배열을 정수형 배열로 변환(빈 문자열은 무시됨)
        for row in pos:
            if row[0] and row[1]:  # 문자열이 비어 있지 않으면
                try:
                    pos_int.append([int(row[0]), int(row[1])])
                except ValueError:
                    print(f"Error converting position to integer: {row}")
                    continue
        
        pos_int = np.array(pos_int)
        norm_speed_x = self.parameter['norm_speed']
        norm_speed_y = self.parameter['norm_speed']

        
        for i, (x, y) in enumerate(pos_int):  # 반복문을 통해 좌표 이동 후 Zero Scan 수행
            self.afm.moveto_xy_stage(x, y, norm_speed_x, norm_speed_y)

            reply = self.afm.query_stage_pos()
            sw_x = reply['x']
            sw_y = reply['y']
            sleep(2)
            
            # 로그 기록
            line = f'{self.iter}, P{i}, {x}, {y}, {sw_x}, {sw_y}\n'
            self.file_name = f'RAW_Repeat_{self.iter}_P{i}'
            self.afm.set_data_location(self.base_dir, self.sub_dir, self.file_name)
            with open(self.log_path, "a") as log_fd:
                log_fd.write(line)


            self.__approach()
            self.__run()
            self.afm.lift_z_stage(1000)



    def done(self):
        self.afm.stop_scan()
        self.afm.stop_scan()
        if self.parameter['lift_enable'] == 'On':
            self.afm.lift_z_stage(1000)

    def stop(self):
        checker = CHECKER()
        checker.abort_approach()
        checker.abort_scan()


    def analyze(self):
        find_dir = self.base_dir + '\\' + self.sub_dir
        tiff_list = src.core.report.ls_tiff(find_dir)
        read_tiff = src.core.tiff.Read_tiff()
        write_tiff = src.core.tiff.Write_tiff()
        
        result_mean = []
        result_min = []
        result_max = []
        result_name = []
        result_rq = []
        for tiff_path in tiff_list:
            tiff_name = tiff_path.split('\\')[-1]
            tmp_name = tiff_name.split('.')[0]
            condition_1 = tmp_name.split('_')[-3] == 'Z Height'
            condition_2 = tmp_name.split('_')[-2] == 'Forward'
            tmp_name_list = tiff_name.split('_')
            tmp_name = ''
            for s in tmp_name_list[1:]: tmp_name += '_' + s
            tmp_name = tmp_name.split('.')[0]
            tmp_name = tmp_name[1:]
            if condition_1 & condition_2:
                dict_tiff , ori_tiff = read_tiff(tiff_path)
                image = dict_tiff['IMAGE']
                flatten_image = src.core.analysis.one_d_flatten(image)
                dict_tiff['IMAGE'] = flatten_image
                cmap = dict_tiff['COLORMAP']
                save_path = self.base_dir + '\\' +self.sub_dir + '\\Flatten\\' + 'Flatten_' + tmp_name + '.tiff'
                report_dir = self.base_dir + '\\' +self.sub_dir + '\\Report\\'
                write_tiff(ori_tiff,dict_tiff,save_path)
                # convert scale : um ->nm
                scale_factor = 1E+3
                tmp_mean, tmp_min, tmp_max = src.core.analysis.cal_statistic(flatten_image)
                tmp_mean *= scale_factor; tmp_min *= scale_factor; tmp_max *= scale_factor
                tmp_rq = src.core.analysis.cal_r_q(flatten_image) * scale_factor
                
                tmp_data = [
                            ['Mean', 'nm',f'{tmp_mean:.3f}'],
                            ['Min', 'nm', f'{tmp_min:.3f}'],
                            ['Max', 'nm', f'{tmp_max:.3f}'],
                            ['Rq', 'nm', f'{tmp_rq:.3f}']
                            ]
                flatten_image *= scale_factor
                src.core.report.default_report(tmp_name, report_dir, dict_tiff, cmap, tmp_data)

                result_name.append(tmp_name)
                result_mean.append(tmp_mean)
                result_min.append(tmp_min)
                result_max.append(tmp_max)
                result_rq.append(tmp_rq)
                
        
        csv_path = self.base_dir + '\\' + self.sub_dir + "\\result.csv"
        result_dict = {'File name': result_name,
                       'Mean (nm)': result_mean,
                       'Min (nm)' : result_min,
                       'Max (nm)' : result_max,
                       'Rq (nm)' :  result_rq
                      } 
        src.core.report.write_csv(result_dict,csv_path,['Mean (nm)','Min (nm)','Max (nm)','Rq (nm)'])

from src.core.afm import AFM as AFM
from src.core.afm import CHECKER as CHECKER
import src.core.analysis
import src.core.report
import src.core.tiff 
from time import sleep
import numpy as np
import os


class Wafer_Zero_Scan():
    
    def __init__(self):
        self.afm = AFM()
        self.base_dir = ''
        self.sub_dir = 'Wafer Zero Scan'
        self.file_name = 'RAW'
        self.parameter = None
    
    def move_to(self):
        pass


    def set(self):
        if self.parameter['xy_servo_enable'] == 'On':
            self.sub_dir += '/XYservo_On'
            self.parameter['z_scanner_range'] = 70
            self.parameter['xy_scanner_range'] = 100
            self.afm.enable_xy_servo()
       
        elif self.parameter['xy_servo_enable'] == 'Off':
            self.sub_dir += '/XYservo_Off'
            self.log_name = 'Position_XY.csv' 
            self.parameter['z_scanner_range'] = 20
            self.parameter['xy_scanner_range'] = 20
            self.afm.enable_xy_servo('off')
            self.header = 'Repeat, Lable , Input X Pos., Input Y Pos., Output X Pos., Output Y Pos.\n '

        if self.parameter['lift_enable'] == 'On':
            self.sub_dir += '_Lift_On'
        elif self.parameter['lift_enable'] == 'Off':
            self.sub_dir += '_Lift_Off'

        dir_path = self.base_dir + '\\' + self.sub_dir
        self.log_path = dir_path + '\\' + self.log_name
        if not os.path.exists(dir_path):os.makedirs(dir_path)
        fd = open(self.log_path,'w')
        fd.write(self.header)
        fd.close()


        self.afm.set_data_location(self.base_dir, self.sub_dir, self.file_name)
        self.afm.set_head_mode('contact')
        self.afm.clear_channels()
        self.afm.add_channel('ChannelZDriveOrTopography')
        
        self.afm.stop_scan()
        self.afm.set_z_scanner_range(self.parameter['z_scanner_range'])
        self.afm.set_xy_scanner_range(self.parameter['xy_scanner_range'])
        self.bck_scan_geometry = self.afm.get_scan_geometry_all()
        self.bck_scan_option = self.afm.get_scan_option_all()
        self.bck_approach_option = self.afm.get_approach_option_all()
        self.afm.set_scan_geometry(self.parameter['pixels'][0],self.parameter['pixels'][1], 0, 0)
        self.afm.set_scan_option(sine_scan=False,over_scan=False, \
                        over_scan_percent=0, tow_way=False, \
                        det_driven=False, force_slope_correction=False, \
                        interlace=False,slow_scan='off')
        self.afm.set_scan_rate(self.parameter['scan_rate'])
        self.afm.enable_z_servo()
        self.afm.set_z_servo(self.parameter['gain'], self.parameter['setpoint'])
        self.afm.set_approach_option()
    


    def unset(self):    
        self.afm.set_scan_option_dict(self.bck_scan_option)
        self.afm.set_scan_geometry_dict(self.bck_scan_geometry)
        self.afm.set_scan_option_dict(self.bck_scan_option)
    
    def __approach(self):
        self.afm.start_approach('q+s')
        wait_time = self.parameter['post_approach_wait']
        print(f"Waiting for {wait_time} seconds")
        sleep(wait_time)
    
    def __run(self):
        self.afm.trigger_image_scan()
        isscan = True
        while isscan:
            test = self.afm.query_scan_status()
            if test == 'true':isscan=True;
            else:isscan=False
            sleep(2)

    def approach(self):
        pass



    def run(self):
        pos = np.array(self.parameter['target_positions'])  # 'target_positions'에서 좌표 데이터 로드
        
        pos_int = []   # 문자열 배열을 정수형 배열로 변환(빈 문자열은 무시됨)
        for row in pos:
            if row[0] and row[1]:  # 문자열이 비어 있지 않으면
                try:
                    pos_int.append([int(row[0]), int(row[1])])
                except ValueError:
                    print(f"Error converting position to integer: {row}")
                    continue
        
        pos_int = np.array(pos_int)
        norm_speed_x = self.parameter['norm_speed']
        norm_speed_y = self.parameter['norm_speed']

        
        for i, (x, y) in enumerate(pos_int):  # 반복문을 통해 좌표 이동 후 Zero Scan 수행
            self.afm.moveto_xy_stage(x, y, norm_speed_x, norm_speed_y)

            reply = self.afm.query_stage_pos()
            sw_x = reply['x']
            sw_y = reply['y']
            sleep(2)
            
            # 로그 기록
            line = f'{self.iter}, P{i}, {x}, {y}, {sw_x}, {sw_y}\n'
            self.file_name = f'RAW_Repeat_{self.iter}_P{i}'
            self.afm.set_data_location(self.base_dir, self.sub_dir, self.file_name)
            with open(self.log_path, "a") as log_fd:
                log_fd.write(line)


            self.__approach()
            self.__run()
            self.afm.lift_z_stage(1000)



    def done(self):
        self.afm.stop_scan()
        self.afm.stop_scan()
        if self.parameter['lift_enable'] == 'On':
            self.afm.lift_z_stage(1000)

    def stop(self):
        checker = CHECKER()
        checker.abort_approach()
        checker.abort_scan()


    def analyze(self):
        find_dir = self.base_dir + '\\' + self.sub_dir
        tiff_list = src.core.report.ls_tiff(find_dir)
        read_tiff = src.core.tiff.Read_tiff()
        write_tiff = src.core.tiff.Write_tiff()
        
        result_mean = []
        result_min = []
        result_max = []
        result_name = []
        result_rq = []
        for tiff_path in tiff_list:
            tiff_name = tiff_path.split('\\')[-1]
            tmp_name = tiff_name.split('.')[0]
            condition_1 = tmp_name.split('_')[-3] == 'Z Height'
            condition_2 = tmp_name.split('_')[-2] == 'Forward'
            tmp_name_list = tiff_name.split('_')
            tmp_name = ''
            for s in tmp_name_list[1:]: tmp_name += '_' + s
            tmp_name = tmp_name.split('.')[0]
            tmp_name = tmp_name[1:]
            if condition_1 & condition_2:
                dict_tiff , ori_tiff = read_tiff(tiff_path)
                image = dict_tiff['IMAGE']
                flatten_image = src.core.analysis.one_d_flatten(image)
                dict_tiff['IMAGE'] = flatten_image
                cmap = dict_tiff['COLORMAP']
                save_path = self.base_dir + '\\' +self.sub_dir + '\\Flatten\\' + 'Flatten_' + tmp_name + '.tiff'
                report_dir = self.base_dir + '\\' +self.sub_dir + '\\Report\\'
                write_tiff(ori_tiff,dict_tiff,save_path)
                # convert scale : um ->nm
                scale_factor = 1E+3
                tmp_mean, tmp_min, tmp_max = src.core.analysis.cal_statistic(flatten_image)
                tmp_mean *= scale_factor; tmp_min *= scale_factor; tmp_max *= scale_factor
                tmp_rq = src.core.analysis.cal_r_q(flatten_image) * scale_factor
                
                tmp_data = [
                            ['Mean', 'nm',f'{tmp_mean:.3f}'],
                            ['Min', 'nm', f'{tmp_min:.3f}'],
                            ['Max', 'nm', f'{tmp_max:.3f}'],
                            ['Rq', 'nm', f'{tmp_rq:.3f}']
                            ]
                flatten_image *= scale_factor
                src.core.report.default_report(tmp_name, report_dir, dict_tiff, cmap, tmp_data)

                result_name.append(tmp_name)
                result_mean.append(tmp_mean)
                result_min.append(tmp_min)
                result_max.append(tmp_max)
                result_rq.append(tmp_rq)
                
        
        csv_path = self.base_dir + '\\' + self.sub_dir + "\\result.csv"
        result_dict = {'File name': result_name,
                       'Mean (nm)': result_mean,
                       'Min (nm)' : result_min,
                       'Max (nm)' : result_max,
                       'Rq (nm)' :  result_rq
                      } 
        src.core.report.write_csv(result_dict,csv_path,['Mean (nm)','Min (nm)','Max (nm)','Rq (nm)'])