from multiprocessing import Condition
from unittest import TestCase
from transitions import Machine, State
from archemist.state.station import Station, StationState
from archemist.state.stations.light_box_station import VialProcessingOpDescriptor,ColourDescriptor
from archemist.state.robot import RobotTaskType
from archemist.state.robots.kukaLBRIIWA import KukaLBRTask
from archemist.util import Location
import archemist.persistence.objectConstructor

class LightBoxSM():
    
    def __init__(self, station: Station, args: dict):

        self._station = station
        self._current_rack_index = 0
        self._processed_racks = 0
        self.operation_complete = False
        self.batch_mode = args['batch_mode']
        #self.rack_load_task = args['rack_load_task']
        #self.rack_unload_task = args['rack_unload_task']
        #self.vial_load_task = args['vial_load_task']
        #self.vial_unload_task = args['vial_unload_task']


        ''' States '''
        states = [State(name='init_state', on_enter='_print_state'), 
            State(name='load_sample', on_enter=['request_load_vial_job','_print_state']),
            State(name='unload_sample', on_enter=['request_unload_vial_job','_print_state']),
            State(name='station_process', on_enter=['request_process_data_job', '_print_state']),
            State(name='update_rack_index', on_enter=['request_rack_index_update', '_print_state']),
            State(name='final_state', on_enter='finalize_batch_processing')]
            
            
        self.machine = Machine(self, states=states, initial='init_state')

        ''' Transitions '''

        self.machine.add_transition('process_state_transitions',source='init_state',dest='load_sample', conditions='all_batches_assigned')

        # load_sample transitions
        self.machine.add_transition('process_state_transitions', source='load_sample',dest='station_process', conditions='is_station_job_ready')

        # station_process transitions
        self.machine.add_transition('process_state_transitions', source='station_process',dest='unload_sample', conditions='is_station_job_ready', before='process_sample')

        # unload_sample transitions
        self.machine.add_transition('process_state_transitions', source='unload_sample',dest='load_sample', conditions='is_station_job_ready', unless='are_all_samples_loaded')
        self.machine.add_transition('process_state_transitions', source='unload_sample',dest='update_rack_index', conditions=['are_all_samples_loaded','is_station_job_ready'], before='reset_station')
        
        self.machine.add_transition('process_state_transitions', source='update_rack_index',dest='load_sample', conditions='is_station_job_ready', unless='are_all_racks_processed')

        self.machine.add_transition('process_state_transitions', source='update_rack_index',dest='final_state', conditions=['are_all_racks_processed','is_station_job_ready'])


    def is_station_job_ready(self):
        return not self._station.has_station_op() and not self._station.has_robot_job()

    def are_all_samples_loaded(self):
        return self._station.loaded_samples == self._station.assigned_batches[self._current_rack_index].num_samples

    def are_all_racks_processed(self):
        return self._processed_racks == self._station.batch_capacity

    def all_batches_assigned(self):
        return not self._station.has_free_batch_capacity()

    def reset_station(self):
        while(self._station.loaded_samples > 0):
          self._station.unload_sample()

    def request_load_vial_job(self):
        self._station.load_sample()
        sample_index = self._station.loaded_samples
        perform_6p = False # this will be later evaluated by the KMRiiwa handler
        allow_auto_func = False # to stop auto charing and calibration when presentig the vial
        robot_job = KukaLBRTask('PresentVial',[perform_6p,self._current_rack_index+1,sample_index,allow_auto_func],RobotTaskType.MANIPULATION, self._station.location)
        current_batch_id = self._station.assigned_batches[self._current_rack_index].id
        self._station.set_robot_job(robot_job,current_batch_id)

    def request_unload_vial_job(self):
        sample_index = self._station.loaded_samples
        perform_6p = False
        if self.are_all_samples_loaded() and (self._processed_racks - 1) == self._station.batch_capacity:
            allow_auto_func = True # enable auto charge and calibration since we are done un/loading samples
        else:
            allow_auto_func = False
        robot_job = KukaLBRTask('ReturnVial',[perform_6p,self._current_rack_index+1,sample_index,allow_auto_func],RobotTaskType.MANIPULATION, self._station.location)
        current_batch_id = self._station.assigned_batches[self._current_rack_index].id
        self._station.set_robot_job(robot_job,current_batch_id)

    def request_rack_index_update(self):
        self._current_rack_index += 1
        self._processed_racks += 1

    def finalize_batch_processing(self):
        self._station.process_assigned_batches()
        self._current_rack_index = 0
        self._processed_racks = 0
        self.to_init_state()
        
    def request_process_data_job(self):
        self._station.set_station_op(VialProcessingOpDescriptor(dict(), ColourDescriptor()))

    def process_sample(self):
        last_operation_op = self._station.station_op_history[-1]
        self._station.assigned_batches[self._current_rack_index].add_station_op_to_current_sample(last_operation_op)
        self._station.assigned_batches[self._current_rack_index].process_current_sample()

    def _print_state(self):
        print(f'[{self.__class__.__name__}]: current state is {self.state}')


