from PymoNNto.Exploration.UI_Base import *
from PymoNNto.Exploration.StorageManager.StorageManager import *
from PymoNNto.Exploration.Evolution.SSH_Functions import *
from PymoNNto.Exploration.Evolution.EvolutionPlots import *
from multiprocessing import Process, Queue, Pipe
import shutil

class QHLine(QFrame):
    def __init__(self):
        super(QHLine, self).__init__()
        self.setFrameShape(QFrame.HLine)
        self.setFrameShadow(QFrame.Sunken)

class QVLine(QFrame):
    def __init__(self):
        super(QVLine, self).__init__()
        self.setFrameShape(QFrame.VLine)
        self.setFrameShadow(QFrame.Sunken)

def evolution_thread_worker(name, python_cmd, conn):
    print('local thread started')
    try:
        dir = os.getcwd()
        os.chdir(get_epc_folder()+'/'+name)

        system_str = sys.platform
        if 'darwin' in system_str:#mac os
            os.system(python_cmd+' execute_evolution.py')#todo: test dfdsgdf
        elif 'win' in system_str:#windows
            os.system('start /wait cmd /c '+python_cmd+' execute_evolution.py')
        else:
            os.system('gnome-terminal -x '+python_cmd+' execute_evolution.py')

        os.chdir(dir)

        conn.send('execution finished')
    except:
        conn.send('thread error')
    conn.send('idle')


class UI_Evolution_Manager(UI_Base):

    def __init__(self):
        super().__init__(None, label='Evolution Monitor', create_sidebar=True)

        self.tabs.currentChanged.connect(self.On_Tab_Changed)

        ssm = SimpleStorageManager(get_epc_folder()+'/')
        servers = eval(ssm.load_param('servers', default="[]"))#['local', 'ssh vieth@poppy.fias.uni-frankfurt.de', 'ssh marius@hey3kmuagjunsk2b.myfritz.net', '+']
        self.listwidget2 = QListWidget()
        self.listwidget2.addItems(['+', 'local']+servers)
        self.listwidget2.currentItemChanged.connect(self.on_server_select)
        self.Add_element(self.listwidget2, stretch=10.0, sidebar=True)

        horizontal_layout = QHBoxLayout()
        self.sidebar_current_vertical_layout.addLayout(horizontal_layout)

        vertical_layout = QVBoxLayout()
        horizontal_layout.addLayout(vertical_layout)
        vertical_layout.addWidget(QLabel('evo name'))
        vertical_layout.addWidget(QLabel('slave file'))
        vertical_layout.addWidget(QLabel('thread number'))
        vertical_layout.addWidget(QLabel('python cmd'))
        vertical_layout.addWidget(QLabel('individual count'))
        vertical_layout.addWidget(QLabel('mutation'))
        vertical_layout.addWidget(QLabel('death rate'))
        vertical_layout.addWidget(QLabel('start geneomes'))
        vertical_layout.addWidget(QLabel('inactive info'))
        vertical_layout.addWidget(QLabel('constriants'))
        vertical_layout.addWidget(QLabel('evo options'))

        self.evo_name_edit = QLineEdit('evo_name')
        self.slave_file_edit = QLineEdit('Exploration/Evolution/test_slave.py')
        self.thread_number_edit = QLineEdit('4')
        self.python_cmd_edit = QLineEdit("python3")
        self.individual_count_edit = QLineEdit('10')
        self.mutation_edit = QLineEdit('0.05')
        self.death_rate_edit = QLineEdit('0.5')
        self.start_genomes_edit = QLineEdit("[{'a':1,'b':2,'c':2,'d':2,'e':3}]")
        self.inactive_genome_info_edit = QLineEdit('{}')
        self.constraints_edit = QLineEdit("['a<b','b<=c']")
        self.evo_options_edit = QLineEdit('{}')


        vertical_layout = QVBoxLayout()
        horizontal_layout.addLayout(vertical_layout)
        vertical_layout.addWidget(self.evo_name_edit)
        vertical_layout.addWidget(self.slave_file_edit)
        vertical_layout.addWidget(self.thread_number_edit)
        vertical_layout.addWidget(self.python_cmd_edit)
        vertical_layout.addWidget(self.individual_count_edit)
        vertical_layout.addWidget(self.mutation_edit)
        vertical_layout.addWidget(self.death_rate_edit)
        vertical_layout.addWidget(self.start_genomes_edit)
        vertical_layout.addWidget(self.inactive_genome_info_edit)
        vertical_layout.addWidget(self.constraints_edit)
        vertical_layout.addWidget(self.evo_options_edit)

        self.add_btn = self.Add_element(QPushButton('Add'), sidebar=True, stretch=1.0)
        self.add_btn.clicked.connect(self.on_add_click)

        for dir in os.listdir(get_epc_folder()):
            if os.path.isdir(get_epc_folder()+'/'+dir):
                self.add_tab(dir)

    def keyPressEvent(self, event):
        if event.key() in [Qt.Key_I]:
            tab = self.tabs.currentWidget()
            print(tab.clicked_id)
            smg = StorageManagerGroup(tab.name, data_folder=get_data_folder() + '/Evolution_Project_Clones/' + tab.name + '/Data')
            smg.sort_by('gen')
            sm = smg.StorageManagerList[tab.clicked_id]
            file = sm.absolute_path + sm.config_file_name
            print(file)
            with open(file) as f:
                lines = f.readlines()

                layout = QVBoxLayout()
                pte = QPlainTextEdit()
                pte.setPlainText('\r\n'.join(lines))
                pte.setReadOnly(True)
                layout.addWidget(pte)
                dlg = QDialog()
                dlg.setWindowTitle("Info")
                dlg.setLayout(layout)
                dlg.resize(600, 400)
                dlg.exec()

    def set_text(self, ssm, edit, text):
        data = ssm.load_param(text, default='', return_string=True)
        if data is not None and data != '':
            edit.setText(data)

    def On_Tab_Changed(self, i):
        if i is not None:
            tab = self.tabs.currentWidget()
            if hasattr(tab, 'ssm'):
                self.update_status(tab, tab.name)
                ssm = self.tabs.currentWidget().ssm
                self.set_text(ssm, self.evo_name_edit, 'name')
                self.set_text(ssm, self.slave_file_edit, 'slave_file')
                self.set_text(ssm, self.inactive_genome_info_edit, 'inactive_genome_info')
                self.set_text(ssm, self.thread_number_edit, 'thread_number')
                self.set_text(ssm, self.individual_count_edit, 'individual_count')
                self.set_text(ssm, self.mutation_edit, 'mutation')
                self.set_text(ssm, self.death_rate_edit, 'death_rate')
                self.set_text(ssm, self.start_genomes_edit, 'start_genomes')
                self.set_text(ssm, self.constraints_edit, 'constraints')
                self.set_text(ssm, self.python_cmd_edit, 'python_cmd')
                self.set_text(ssm, self.evo_options_edit, 'evo_options')


    def update_status(self, tab, name):
        if 'ssh ' in tab.server:
            if name in ssh_get_running(tab.server):
                tab.status_label.setText('<font color="green">running</font>')
                return True
            else:
                tab.status_label.setText('<font color="red">inactive</font>')
                return False

        if 'local' in tab.server:
            if tab.process is not None:
                if tab.parent_conn.poll():
                    msg = tab.parent_conn.recv()
                    tab.process.kill()
                    tab.process = None
                    tab.status_label.setText('<font color="red">inactive</font>')
                    return False
                else:
                    tab.status_label.setText('<font color="green">running</font>')
                    return True
            else:
                tab.status_label.setText('<font color="red">inactive</font>')
                return False

    def add_tab(self, name):

        folder = get_epc_folder()+'/'+name+'/'

        tab = self.Next_Tab(name, stretch=0.0)
        tab.name=name

        tab.status_label = self.Add_element(QLabel('Status'))

        tab.ssm = SimpleStorageManager(folder)
        tab.server = tab.ssm.load_param('server', default=None)

        tab.gene_keys = tab.ssm.load_param('gene_keys', default=None)
        if tab.gene_keys is not None:
            tab.gene_keys = eval(tab.gene_keys)#should return list like ['a', 'b', 'c', 'd', 'e']  todo:save on evo start

        def refresh():
            try:
                print('refresh')
                if tab.server is not None and 'ssh ' in tab.server:
                    user, host, password = split_ssh_user_host_password_string(tab.server)
                    print('get Data')
                    get_Data(name, user, host, password)
                if tab.gene_keys is not None:
                    print('update plot')
                    update_evolution_plot(self, tab, name, tab.gene_keys, data_folder=get_data_folder() + '/Evolution_Project_Clones/' + name + '/Data')
            except:
                print('no plot data found in', name)

            self.update_status(tab, name)

        refresh_btn = self.Add_element(QPushButton('refresh state'))
        refresh_btn.clicked.connect(refresh)

        tab.process = None

        tab.python_cmd = tab.ssm.load_param('python_cmd', default='python3')

        def start_continue():
            if 'ssh ' in tab.server:
                if not self.update_status(tab, name):
                    ssh_execute_evo(tab.server, name, python_cmd=tab.python_cmd)#python3
                else:
                    print('thread already running')

            if 'local' in tab.server:
                if tab.process is None:
                    tab.parent_conn, child_conn = Pipe()
                    tab.process = Process(target=evolution_thread_worker, args=(name, tab.python_cmd, child_conn))#python
                    tab.process.start()
                else:
                    print('thread already running')

            self.update_status(tab, name)

        start_evo_btn = self.Add_element(QPushButton('start/continue'))
        start_evo_btn.clicked.connect(start_continue)

        def stop():
            if 'ssh ' in tab.server:
                ssh_stop_evo(tab.server, name)
                self.update_status(tab, name)
            if 'local' in tab.server:
                if tab.process is not None:
                    print('Please close terminal manually')
                    tab.process.kill()
                    tab.process.terminate()
                    tab.process.close()
                    tab.process=None


        stop_btn = self.Add_element(QPushButton('stop evolution'))
        stop_btn.clicked.connect(stop)


        def remove():
            if 'ssh ' in tab.server:
                ssh_stop_evo(tab.server, name, remove_evo=True)

            if 'local' in tab.server:
                if tab.process is not None:
                    print('Please close terminal manually')
                    tab.process.kill()
                    tab.process.terminate()
                    tab.process.close()
                    tab.process=None

            shutil.rmtree(get_epc_folder() + '/' + name + '/')
            self.tabs.removeTab(self.tabs.currentIndex())

        remove_btn = self.Add_element(QPushButton('remove'))
        remove_btn.clicked.connect(remove)

        def help():
            layout = QVBoxLayout()
            pte = QPlainTextEdit()
            pte.setPlainText('''Every tab corresponds to a folder in:
Data/Evolution_Project_Clones/
            
For remote SSH servers the local folder only contains some config files while the project copy is located in the user folder of the server.

When the evolution is started on a ssh server it creates a "screen" session which is named after the evolution name.

Useful functions are:

"screen -list" to show the screen sessions and their ids,
"screen -r id" to connect to a session where "id" has to be replaced and
"exit" to close a screen session you are connected to.

The project folders also contain the file execute_evolution.py which is used to start the evolution.

All genomes generate a file with the results and the genes when set_score is called in the slave file. Which are transferred to the local folder when the state is refreshed.
''')
            pte.setReadOnly(True)
            layout.addWidget(pte)

            #self.layout.addWidget(message)
            #self.layout.addWidget(self.buttonBox)

            dlg = QDialog()
            dlg.setWindowTitle("Help")
            dlg.setLayout(layout)
            dlg.resize(600, 400)
            dlg.exec()

        help_btn = self.Add_element(QPushButton('help'))
        help_btn.clicked.connect(help)

        self.Add_element(QLabel(tab.server))

        self.Next_H_Block(stretch=10)

        add_evolution_plot_items(self, tab)

        tab.folder = get_epc_folder() + '/' + name + '/'

        #self.update_status(tab, name)
        self.On_Tab_Changed(None)

    def on_server_select(self):
        server_str = self.listwidget2.currentItem().text()
        if server_str == '+':
            text, ok = QInputDialog.getText(self.main_window, 'Add...', 'Add computing device "ssh user@host" or "ssh user@host password"')
            if ok and text != '+':
                self.listwidget2.addItem(text)
                ssm = SimpleStorageManager(get_epc_folder()+'/')
                print(ssm.absolute_path)
                ssm.save_param('servers', eval(ssm.load_param('servers', default="[]"))+[text])

    def on_file_click(self):
        print(self.listwidget.currentItem().text())

    def on_add_click(self):
        name = self.evo_name_edit.text()
        if is_invalid_evo_name(name):
            print('invalid evolution name.\r\n For savety reasons some names and characters are forbidden to avoid the accidental removal of files or folders')
        else:
            item = self.listwidget2.currentItem()
            if item is None:
                print('please select a server first')
            else:
                server_str = item.text()
                if os.path.exists(get_data_folder()+'/Evolution_Project_Clones/'+name):
                    print('evolution folder already exists')
                else:

                    valid_genomes = False
                    try:
                        sg=eval(self.start_genomes_edit.text())
                        for key in ['gen', 'score', 'id']:
                            for genome in sg:
                                if key in genome:
                                    genome.pop(key)
                                    print(key, 'removed from genomes')
                        self.start_genomes_edit.setText(str(sg))
                        valid_genomes = True
                    except:
                        print('error parsing start genomes')

                    if not valid_genomes:
                        print('error parsing start genomes')
                    else:

                        self.copy_project(name, server_str)

                        folder = get_epc_folder()+'/'+name+'/'
                        ssm = SimpleStorageManager(folder)
                        ssm.save_param('server', server_str)

                        ssm.save_param('name', name)
                        ssm.save_param('slave_file', self.slave_file_edit.text())
                        ssm.save_param('inactive_genome_info', self.inactive_genome_info_edit.text())
                        ssm.save_param('thread_number', self.thread_number_edit.text())
                        ssm.save_param('individual_count', self.individual_count_edit.text())
                        ssm.save_param('mutation', self.mutation_edit.text())
                        ssm.save_param('death_rate', self.death_rate_edit.text())
                        ssm.save_param('start_genomes', self.start_genomes_edit.text())
                        ssm.save_param('constraints', self.constraints_edit.text())
                        ssm.save_param('python_cmd', self.python_cmd_edit.text())
                        ssm.save_param('evo_options', self.evo_options_edit.text())

                        exec_file = """
from PymoNNto.Exploration.Evolution.Evolution import *
if __name__ == '__main__':
    evo = Evolution(name='#name#',
                slave_file='#slave_file#',
                individual_count=#individual_count#,
                mutation=#mutation#,
                death_rate=#death_rate#,
                constraints=#constraints#,
                inactive_genome_info=#inactive_genome_info#,
                start_genomes=#start_genomes#,
                devices={'multi_thread': #thread_number#},
                additional_evo_params=#evo_options#
                )
    
    if not evo.start(ui=False):
        evo.continue_evolution(ui=False)
"""

                        print('generate execute_evolution.py...')
                        exec_file = exec_file.replace('#name#', name)
                        exec_file = exec_file.replace('#slave_file#', self.slave_file_edit.text())
                        exec_file = exec_file.replace('#inactive_genome_info#', self.inactive_genome_info_edit.text())
                        exec_file = exec_file.replace('#thread_number#', self.thread_number_edit.text())
                        exec_file = exec_file.replace('#individual_count#', self.individual_count_edit.text())
                        exec_file = exec_file.replace('#mutation#', self.mutation_edit.text())
                        exec_file = exec_file.replace('#death_rate#', self.death_rate_edit.text())
                        exec_file = exec_file.replace('#start_genomes#', self.start_genomes_edit.text())
                        exec_file = exec_file.replace('#constraints#', self.constraints_edit.text())
                        exec_file = exec_file.replace('#evo_options#', self.evo_options_edit.text())

                        evo_file = folder + 'execute_evolution.py'

                        md_file = open(folder + 'execute_evolution.py', "w")
                        md_file.write(exec_file)
                        md_file.close()

                        print('saving evolution config...')

                        gene_keys = list(eval(self.start_genomes_edit.text())[0].keys())
                        ssm.save_param('gene_keys', gene_keys)

                        if 'ssh' in server_str:
                            user, host, password = split_ssh_user_host_password_string(server_str)
                            print('transfer execute_evolution.py to ' + host + '...')
                            ssh = get_ssh_connection(host, user, password)
                            scp = SCPClient(ssh.get_transport())
                            scp.put(evo_file, name + '/execute_evolution.py')
                            scp.close()
                            ssh.close()

                        print('adding tab...')
                        self.add_tab(name)

    def copy_project(self, evo_name, device_str):
        if device_str == 'local':
            clone_project(evo_name)

        if 'ssh' in device_str:
            user, host, password = split_ssh_user_host_password_string(device_str)
            if user is not None and host is not None:
                transfer_project(evo_name, user, host, password)

            Data_Folder = get_data_folder()
            if Data_Folder != './Data':
                epc_folder = Data_Folder + '/Evolution_Project_Clones'
                create_folder_if_not_exist(epc_folder)
                run_folder = epc_folder + '/' + evo_name
                create_folder_if_not_exist(run_folder)


if __name__ == '__main__':
    UI_Evolution_Manager().show()

