# -*- coding: utf-8 -*-
# __author__ = 'sanger'

from __future__ import division
import os

import math

from biocluster.config import Config
from bson.objectid import ObjectId
import types
import json
import numpy as np
import re
from types import StringTypes
import gridfs
from collections import OrderedDict, defaultdict
import pandas as pd
from biocluster.file import getsize, exists
from biocluster.file import download
import shutil
from biocluster.api.file.lib.transfer import MultiFileTransfer
import sys

project_type = 'lnc_rna'
db = Config().get_mongo_client(mtype=project_type)[Config().get_mongo_dbname(project_type)]

# ---------------------基因集及相关gdq----------------------------------
def export_geneset_exp_matrix(data, option_name, dir_path, bind_obj=None):
    """ gdq 2019 04 修改
    该函数根据提供的geneset_id和exp_id和group_dict提取表达信息。
    当geneset_id为‘all’时，则导出所有的exp_id对应的详情表。
    当geneset_id为‘refall’时，则导出所有的exp_id对应的详情表时还需限制'is_new'字段为False(false)。
    该函数还顺便根据group_dict信息生成分组信息文件group_info.txt。
    该函数仅返回表达矩阵文件的路径信息
    """
    exp_id, geneset_id, geneset_type = data.split(";")
    group_dict = bind_obj.sheet.option('group_dict')
    group_dict = json.loads(group_dict, object_pairs_hook=OrderedDict)
    samples = list()
    for each in group_dict:
        samples += group_dict[each]
    # export group info
    with open(dir_path + '/group_info.txt', 'w') as f:
        f.write('#sample\tgroup\n')
        group_id = bind_obj.sheet.option('group_id').lower()
        for g in group_dict:
            for s in group_dict[g]:
                if group_id == "all":
                    g = s
                f.write('{}\t{}\n'.format(s, g))
    # decide output columns
    target_cols = OrderedDict(seq_id=1, _id=0)
    for each in samples:
        target_cols[each] = 1
    # get geneset
    seq_marker = {}
    geneset = set()
    if "all" not in geneset_id.lower():
        conn = db['sg_geneset_detail']
        for gs_id, gs_type in zip(re.split('\s*,\s*', geneset_id), re.split('\s*,\s*', geneset_type)):
            geneset_record = conn.find({"geneset_id": ObjectId(gs_id)})
            if not geneset_record:
                bind_obj.set_error('geneset not found by query: {}'.format(geneset_id))
            sub_geneset = geneset_record[0]['seq_list']
            geneset.update(sub_geneset)
            for seq_id in sub_geneset:
                seq_marker[seq_id] = gs_type

    # get all exp matrix
    conn = db['sg_exp_detail']
    if "refall" in geneset_id.lower():
        exp_records = conn.find({"exp_id": ObjectId(exp_id), 'is_new': False}, target_cols)
    else:
        exp_records = conn.find({"exp_id": ObjectId(exp_id)}, target_cols)
    exp_matrix = pd.DataFrame(list(exp_records))
    if exp_matrix.shape[0] == 0:
        bind_obj.set_error('No expression data find by query: {}'.format(exp_id))
    exp_matrix = exp_matrix.set_index('seq_id')
    if geneset:
        needed_index = exp_matrix.index & geneset
        exp_matrix = exp_matrix.loc[needed_index, :]
    output = os.path.join(dir_path, option_name)
    exp_matrix.to_csv(output, sep='\t', header=True, index=True)

    marker = os.path.join(dir_path, 'seq_id_marker.json')
    with open(marker, 'w') as out_handler:
        json.dump(seq_marker, out_handler, indent=4)

    print('success to export expression matrix')
    return output

def export_geneset_from_query(data, option_name, dir_path, bind_obj=None):
    collection = db['sg_annotation_query_detail']
    main_collection = db['sg_annotation_query']
    my_result = main_collection.find_one({'task_id': data})
    if not my_result:
        bind_obj.set_error("意外错误，task_id:{}在sg_annotation_query中未找到！".format(data))
    if 'main_id' in my_result:
        query_id = my_result["main_id"]
    else:
        query_id = my_result['_id']
    results = collection.find({"query_id": ObjectId(query_id)})
    output = os.path.join(dir_path, "all_genes_transcripts.txt")
    with open(output, "w") as f:
        f.write('transcript_id\tgene_id\n')
        for result in results:
            transcript_id = ''
            if u'transcript_id' in result:
                transcript_id = result['transcript_id']
            gene_id = ''
            if u'gene_id' in result:
                gene_id = result['gene_id']
            f.write('{}\t{}\n'.format(transcript_id, gene_id))
    return output

def export_wgcna_exp_matrix(data, option_name, dir_path, bind_obj=None):
    """ gdq
    该函数仅为wgcna module分析服务。
    该函数会返回wgcna_prepare得到的exp_matrix路径，
    同时还会通过查询sg_annotation_query表导出geneid和genename的关系文件:seq_id2gene_name.txt
    """
    wgcna_prepare = db['sg_wgcna_prepare']
    prepare_id = data.strip()
    prepare_main = wgcna_prepare.find_one({"_id": ObjectId(prepare_id)})
    exp_matrix = prepare_main['exp_matrix']
    if re.match(r'^\w+://\S+/.+$', exp_matrix):
        download_s3_file(exp_matrix, os.path.join(bind_obj.work_dir, "exp_matrix.txt"))
        exp_matrix = os.path.join(bind_obj.work_dir, "exp_matrix.txt")
        if not os.path.exists(os.path.join(dir_path, "exp_matrix.txt")):
            os.link(exp_matrix, os.path.join(dir_path, "exp_matrix.txt"))
    else:
        if os.path.exists(os.path.join(dir_path, "exp_matrix.txt")):
            os.remove(os.path.join(dir_path, "exp_matrix.txt"))
        os.link(exp_matrix, os.path.join(dir_path, "exp_matrix.txt"))
    target_seqs = pd.read_table(exp_matrix, header=0, index_col=0).index
    task_id = prepare_main['task_id']
    annot_table = db['sg_annotation_query']
    try:
        annot_main = annot_table.find_one({"task_id": task_id, "type": "latest", "status": "end"})
    except:
        bind_obj.set_error("cannot find sg_annotation_query main table")
    else:
        if annot_main is None:
            annot_main = annot_table.find_one({"task_id": task_id, "type": "origin", "status": "end"})
    if "main_id" not in annot_main:
        annot_main_id = annot_main['_id']
    else:
        annot_main_id = annot_main['main_id']
    annot_detail = db['sg_annotation_query_detail']
    exp_level = bind_obj.sheet.option('exp_level')
    if exp_level[0].upper() == 'T':
        query_dict = dict(query_id=annot_main_id, )
        result_dict = dict(_id=0, gene_name=1, gene_id=1, transcript_id=1)
        result = annot_detail.find(query_dict, result_dict)
        gene2name = pd.DataFrame(list(result))
        gene2name.set_index('transcript_id', inplace=True)
    else:
        query_dict = dict(query_id=annot_main_id, is_gene=True)
        result_dict = dict(_id=0, gene_name=1, gene_id=1)
        result = annot_detail.find(query_dict, result_dict)
        gene2name = pd.DataFrame(list(result))
        gene2name.set_index('gene_id', inplace=True)
    gene2name = gene2name.loc[list(target_seqs), :]
    gene2name.reset_index(inplace=True)
    output = os.path.join(dir_path, "seq_id2gene_name.txt")
    gene2name.to_csv(output, sep='\t', header=True, index=False)
    gene2name = pd.read_table(output, header=0)
    gene2name.fillna(method="pad", axis=1, inplace=True)
    gene2name.to_csv(output, sep='\t', header=True, index=False)
    return exp_matrix + ';' + output

# added  by gdq for wgcna
def export_wgcna_relate_input(data, option_name, dir_path, bind_obj=None):
    module_id = data.strip()
    # export eigengenes
    eigengenes = db['sg_wgcna_module_eigengenes_detail']
    eigengenes_found = eigengenes.find({"module_id": ObjectId(module_id)}, {"_id": 0, "module_id": 0})
    eigengenes_pd = pd.DataFrame(list(eigengenes_found))
    eigengenes_pd.set_index("module", inplace=True)
    eigengenes_path = os.path.join(dir_path, "module_eigengenes.xls")
    eigengenes_pd.to_csv(eigengenes_path, sep='\t', header=True, index=True)
    # export exp matrix
    prepare_id = db['sg_wgcna_module'].find_one({"main_id": ObjectId(module_id)})['wgcna_prepare_id']
    exp_matrix = db['sg_wgcna_prepare'].find_one({"main_id": ObjectId(prepare_id)})['exp_matrix']
    # export each gene's module and gene_id/gene_name info
    exp_level = bind_obj.sheet.option('exp_level')
    membership = db['sg_wgcna_module_membership_detail']
    query_dict = dict(module_id=ObjectId(module_id))
    return_dict = dict(_id=0, seq_id=1, gene_name=1, module=1, kme=1, block_id=1)
    if exp_level == 'transcript':
        return_dict.update({"gene_id": 1})
    result = membership.find(query_dict, return_dict)
    result_pd = pd.DataFrame(list(result))
    result_pd.set_index("seq_id", inplace=True)
    gene_annot = os.path.join(dir_path, "seq_annot.xls")
    result_pd.to_csv(gene_annot, sep='\t', header=True, index=True)
    return exp_matrix + ';' + eigengenes_path + ";" + gene_annot

def get_after_qc_bam_path(data, option_name, dir_path, bind_obj=None):
    db = Config().get_mongo_client(mtype="lnc_rna")[Config().get_mongo_dbname("lnc_rna")]
    conn = db['sg_specimen']
    result = conn.find({"task_id": data, "about_qc": "after"})
    sorted_result = sorted(result, key=lambda k: k['_id'])
    output = os.path.join(dir_path, option_name)
    with open(output, 'w') as f:
        for i in sorted_result:
            f.write(i['bam_path'] + "\n")
    return output

# ---------------------从蛋白复制过来的------------------------

def export_compare_exp_fc(data, option_name, dir_path, bind_obj=None):
    '''
    导出差异分析accession_id 和 fc
    '''

    diff_id = bind_obj.sheet.option('diff_id')
    compare_group = bind_obj.sheet.option('compare_group')
    target_cols = OrderedDict(seq_id=1, log2fc=1, _id=0)

    bind_obj.logger.debug("导出表达参数 {}".format(target_cols))
    conn = db['sg_diff_detail']
    diff_exp_records = conn.find({"diff_id": ObjectId(diff_id), "compare": compare_group}, target_cols)
    diff_exp_matrix = pd.DataFrame(list(diff_exp_records))
    output = os.path.join(dir_path, option_name)
    diff_exp_matrix.to_csv(output, sep='\t', header=True, index=False)
    print('success to export expression matrix')
    return output

def export_go_enrich_matrix(data, option_name, dir_path, bind_obj=None):
    go_enrich_id = bind_obj.sheet.option('enrich_id')
    go_type = bind_obj.sheet.option('go_type')

    target_cols = OrderedDict(go_id=1, go_type=1, discription=1, p_corrected=1, p_uncorrected=1, seq_list=1, depth=1,
                              _id=0)
    bind_obj.logger.debug("导出GO富集表")
    # get geneset
    conn = db['sg_geneset_go_enrich_detail']
    if go_type == 'ALL' or go_type == 'All':
        go_enrich_records = conn.find({"go_enrich_id": ObjectId(go_enrich_id)}, target_cols)
    else:
        go_enrich_records = conn.find({"go_enrich_id": ObjectId(go_enrich_id), "go_type": go_type}, target_cols)
    go_enrich_matrix = pd.DataFrame(list(go_enrich_records))
    seq_list2 = [i.replace(';', '|') for i in go_enrich_matrix['seq_list']]
    go_enrich_matrix['seq_list'] = seq_list2
    # exp_matrix = exp_matrix.loc[geneset, :]
    output = os.path.join(dir_path, option_name)
    go_enrich_matrix.to_csv(output, sep='\t', header=True, index=False)
    print('success to export go_enrich matrix')
    return output

def export_kegg_enrich_matrix(data, option_name, dir_path, bind_obj=None):
    '''
    导出kegg富集表格
    '''
    kegg_enrich_id = bind_obj.sheet.option('enrich_id')
    target_cols = OrderedDict(_id=0, id=1, term=1, pvalue=1, corrected_pvalue=1, seq_list=1, kegg_type=1)
    bind_obj.logger.debug("导出KEGG富集表")
    # get geneset
    conn = db['sg_geneset_kegg_enrich_detail']
    bind_obj.logger.debug("导出以下列{}".format(target_cols))
    bind_obj.logger.debug("导出以下列{}".format(kegg_enrich_id))
    kegg_enrich_records = conn.find({"kegg_enrich_id": ObjectId(kegg_enrich_id)}, target_cols)
    kegg_enrich_matrix = pd.DataFrame(list(kegg_enrich_records))
    seq_list2 = [i.replace(';', '|') for i in kegg_enrich_matrix['seq_list']]
    kegg_enrich_matrix['seq_list'] = seq_list2
    # exp_matrix = exp_matrix.loc[geneset, :]
    output = os.path.join(dir_path, option_name)
    kegg_enrich_matrix.to_csv(output, sep='\t', header=True, index=False)
    print('success to export kegg_enrich matrix')
    return output

def export_gene_list_ppi(data, option_name, dir_path, bind_obj=None):
    gene_list_path = os.path.join(dir_path, "%s.txt" % option_name)
    bind_obj.logger.debug("正在导出蛋白集")
    collection = db['sg_geneset_detail']
    main_collection = db['sg_geneset']
    my_result = main_collection.find_one({'_id': ObjectId(data)})
    if not my_result:
        bind_obj.set_error("意外错误，geneset_id:{}在sg_geneset中未找到！".format(ObjectId(data)))
    results = collection.find_one({"geneset_id": ObjectId(data)})["seq_list"]
    with open(gene_list_path, "wb") as f:
        f.write("accession_id" + "\n")
        for result in results:
            f.write(result + "\n")
    bind_obj.logger.debug("蛋白集导出成功！")
    return gene_list_path

def download_s3_file(path, to_path):
    """
    判断文件是否在对象存储上
    """
    if not to_path.startswith("/"):
        to_path = os.path.join(self.work_dir, to_path)
    if os.path.exists(to_path):
        os.remove(to_path)
    elif os.path.exists(path):
        to_path = path
    elif exists(path):
        download(path, to_path)
    else:
        print
        'file can not find'
    return to_path

# added for transcription factor analysis  --------------------------------------------
def get_all_pep_seq(data, option_name, dir_path, bind_obj=None):
    pep_db_path = data.strip()
    result_path = os.path.join(dir_path, "all_pep.fa")
    if re.match(r'^\w+://\S+/.+$', pep_db_path) or re.match(r'/mnt/ilustre', pep_db_path):
        transfer = MultiFileTransfer()
        transfer.add_download(pep_db_path, bind_obj.work_dir + "/")
        transfer.perform()
        # download_s3_file(pep_db_path, os.path.join(bind_obj.work_dir, "pep.db"))
        pep_db_path = os.path.join(bind_obj.work_dir, "refrna_seqs.db")
    bind_obj.logger.debug("文件为 {}".format(pep_db_path))
    import sqlite3
    conn = sqlite3.connect(pep_db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT transcript_id,pep_seq FROM trans_annot")
    with open(result_path, 'w') as fw:
        # pep_id 就是转录本的id
        for pep_id, pep_seq in cursor.fetchall():
            fw.write('>{}\n{}\n'.format(pep_id, pep_seq))
    return result_path

def get_gene_detail(data, option_name, dir_path, bind_obj=None):
    task_id = data.strip()
    annot_table = db['sg_annotation_query']
    annot_main = annot_table.find_one({"task_id": task_id, "type": "latest", 'status': 'end'})
    if not annot_main:
        annot_main = annot_table.find_one({"task_id": task_id, "type": "origin"})
    if not annot_main:
        task_table = db['sg_task']
        task_main = task_table.find_one({"task_id": task_id})
        if os.path.exists(task_main['annot']):
            gene_annot = os.path.join(dir_path, "seq_annot.xls")
            os.system("cut -f 1-4 " + task_main['annot'] + " |awk '{print $2\"\\t\"$1\"\\t\"$4}' > " + gene_annot)
            return gene_annot
            # os.system("cut -f 1-4 {} |awk '{print $2\"\t\"$1\"\t\"$4}' > {} ".format(task_main['annot'], gene_annot))
        else:
            bind_obj.set_error("Not Found in sg_annotation_query by query {}".format(task_id))
    if "main_id" not in annot_main:
        annot_main_id = annot_main['_id']
    else:
        annot_main_id = annot_main['main_id']
    annot_detail = db['sg_annotation_query_detail']
    query_dict = dict(query_id=annot_main_id, )
    result_dict = dict(_id=0, gene_name=1, gene_id=1, description=1, transcript_id=1)
    result = annot_detail.find(query_dict, result_dict)
    result_pd = pd.DataFrame(list(result))
    result_pd.set_index("transcript_id", inplace=True)
    result_pd = result_pd.loc[:, ["gene_id", "gene_name", "description"]]
    result_pd.columns = ["gene_id", "gene_name", "gene_desc"]
    # result_pd = result_pd.loc[:, ["gene_id", "gene_name"]]
    # result_pd.columns = ["gene_id", "gene_name"]
    result_pd["gene_name"][result_pd['gene_name'] == "-"] = None
    result_pd = result_pd.fillna(method="pad", axis=1)
    gene_annot = os.path.join(dir_path, "seq_annot.xls")
    result_pd.to_csv(gene_annot, sep='\t', header=True, index=True)
    os.system(
        r"sed -i 's/%2B/+/g;s/%2F/\//g;s/%2C/,/g;s/%3A/:/g;s/%3B/;/g;s/%3D/=/g;s/%3F/?/g;s/%20/ /g;s/%25/%/g;s/%3C/</g;s/%3E/>/g;s/%5B/[/g;s/%5D/]/g;s/%7B/{/g;s/%7D/}/g' " + gene_annot)
    return gene_annot

# ---------------------基因集及相关zhaozhipeng----------------------------------
def export_genesets_exp_matrix(data, option_name, dir_path, bind_obj=None):
    """
    该函数根据提供的geneset_id和exp_id和group_dict提取表达信息。
    当geneset_id为‘all’时，则导出所有的exp_id对应的详情表。
    当geneset_id为‘refall’时，则导出所有的exp_id对应的详情表时还需限制'is_new'字段为False(false)。
    该函数还顺便根据group_dict信息生成分组信息文件group_info.txt。
    该函数仅返回表达矩阵文件的路径信息
    """
    s_list = data.split(";")
    exp_id = s_list[0]
    geneset_ids = s_list[1:]
    group_dict = bind_obj.sheet.option('group_dict')
    group_dict = json.loads(group_dict, object_pairs_hook=OrderedDict)
    samples = list()
    for each in group_dict:
        samples += group_dict[each]
    # export group info
    with open(dir_path + '/group_info.txt', 'w') as f:
        f.write('#sample\tgroup\n')
        group_id = bind_obj.sheet.option('group_id').lower()
        for g in group_dict:
            for s in group_dict[g]:
                if group_id == "all":
                    g = s
                f.write('{}\t{}\n'.format(s, g))
    # decide output columns
    target_cols = OrderedDict(seq_id=1, _id=0)
    for each in samples:
        target_cols[each] = 1
    # get geneset
    geneset = set()
    bind_obj.logger.debug('=========== %s ++++++++++++++' % ','.join(geneset_ids))
    for geneset_id in geneset_ids:
        if "all" not in geneset_id.lower():
            conn = db['sg_geneset_detail']
            geneset_records = conn.find_one({"geneset_id": ObjectId(geneset_id)})
            if not geneset_records:
                bind_obj.set_error('geneset not found by query: {}'.format(geneset_id))
            geneset.update(geneset_records['seq_list'])

    # get all exp matrix
    conn = db['sg_exp_detail']
    exp_records = conn.find({"exp_id": ObjectId(exp_id)}, target_cols)
    exp_matrix = pd.DataFrame(list(exp_records))
    if exp_matrix.shape[0] == 0:
        bind_obj.set_error('No expression data find by query: {}'.format(exp_id))

    exp_matrix = exp_matrix.set_index('seq_id')
    if geneset:
        indexs = exp_matrix.index & {i for i in geneset}
        exp_matrix = exp_matrix.loc[indexs, :]
    exp_matrix.dropna(how='all')
    output = os.path.join(dir_path, option_name)
    exp_matrix.to_csv(output, sep='\t', header=True, index=True)
    print('success to export expression matrix')
    return output

def export_genes_list(data, option_name, dir_path, bind_obj=None):
    """

    :param data:
    :param option_name:
    :param dir_path:
    :param bind_obj:
    :return:
    """
    bind_obj.logger.debug('=========== %s ++++++++++++++' % data)
    geneset_ids = [item.split(',') for item in data.split(";")]
    geneset = {}
    for gene_type, geneset_id in geneset_ids:
        if "all" not in geneset_id.lower():
            conn = db['sg_geneset_detail']
            geneset_records = conn.find_one({"geneset_id": ObjectId(geneset_id)})
            if not geneset_records:
                bind_obj.set_error('geneset not found by query: {}'.format(geneset_id))
            for i in geneset_records['seq_list']:
                geneset[i] = 'lncRNA' if gene_type == 'l_geneset_id' else 'mRNA'

    output = os.path.join(dir_path, 'genes_info.json')
    with open(output, 'w') as in_handler:
        json.dump(geneset, in_handler, indent=4)

    return output

def export_diff_exp_matrix(data, option_name, dir_path, bind_obj=None):
    """

    :param data: {'diff_group': 'xxx|yyy', 'sg_diff': '5c9da1ba17b2bf206a55d350'}
    :param option_name:
    :param dir_path:
    :param bind_obj:
    :return:
    """
    bind_obj.logger.debug(' :: export diff_exp.txt file')
    compare = data['diff_group']
    sg_diff_id = data['sg_diff_id']

    conn = db['sg_diff_detail']
    recoreds = conn.find({'diff_id': ObjectId(sg_diff_id), 'compare': compare})
    res_data = []
    for dic in recoreds:
        tmp_dic = {'seq_id': dic['seq_id'], 'log2fc': dic['log2fc'], 'regulate': dic['regulate'],
                   'avg_exp': (dic['group1'] + dic['group2']) / 2, 'significant': dic['significant']}
        res_data.append(tmp_dic)
    result_df = pd.DataFrame(res_data)
    output = os.path.join(dir_path, 'diff_exp.txt')
    result_df.to_csv(output, sep='\t', index=False)
    bind_obj.logger.debug(' :: export diff_exp.txt file :: completed')

    return output

def export_rna_type_matrix(data, option_name, dir_path, bind_obj=None):
    """

    :param data: {'task_id': 'lnc_rna', 'exp_level': 'T', 'sg_diff': '5c9da1ba17b2bf206a55d350'}
    :param option_name:
    :param dir_path:
    :param bind_obj:
    :return:
    """
    bind_obj.logger.debug(' :: export seq_id_type.txt file')
    task_id = data['task_id']
    sg_diff_id = data['diff_exp']
    exp_level = data['exp_level']
    # {'task_id': data.task_id, 'exp_level': 'T', 'diff_exp': data.diff_exp}
    conn = db['sg_diff']
    bind_obj.logger.debug(str({'main_id': ObjectId(sg_diff_id), 'exp_level': exp_level, 'task_id': task_id}))
    # recored = conn.find_one({'main_id': ObjectId(sg_diff_id), 'exp_level': exp_level, 'task_id': task_id})
    # 因为main_id本身唯一,所以删除多余查找字段
    recored = conn.find_one({'main_id': ObjectId(sg_diff_id)})
    exp_id = ObjectId(recored['exp_id'])

    conn = db['sg_exp_detail']
    recoreds = conn.find({'exp_id': exp_id}, {'seq_id': 1, 'rna_type': 1, '_id': 0})
    result_df = pd.DataFrame([i for i in recoreds])
    output = os.path.join(dir_path, 'seq_id_type.txt')
    result_df.to_csv(output, sep='\t', index=False)
    bind_obj.logger.debug(' :: export seq_id_type.txt file :: completed')

    return output

def export_cis_matrix2(data, option_name, dir_path, bind_obj=None):
    """

    :param data: {'task_id': 'lnc_rna', 'exp_level': 'T', 'sg_diff': '5c9da1ba17b2bf206a55d350'}
    :param option_name:
    :param dir_path:
    :param bind_obj:
    :return:
    """
    bind_obj.logger.debug(' :: export cis_targets.txt file')

    # 获取基因转录本对应关系， 只要一组

    task_id = bind_obj.sheet.option("task_id")
    annot_table = db['sg_annotation_query']
    annot_main = annot_table.find_one({"task_id": task_id, "type": "origin", "status": "end"})
    annot_detail = db['sg_annotation_query_detail']
    result = annot_detail.find({"query_id": annot_main["main_id"], "is_gene": True},
                               {"_id": 0, "transcript_id": 1, "gene_id": 1})
    g2t_df = pd.DataFrame(list(result))

    conn = db['sg_target_cistrans_detail']
    recoreds = conn.find({'target_cistrans_id': ObjectId(data), "target_type": "cis"},
                         {'lncrna_id': 1, 'gene_id': 1, 'mgene_id': 1, '_id': 0})
    data = [i for i in recoreds]
    output = os.path.join(dir_path, 'cis_targets.txt')
    if data:
        result_df = pd.DataFrame(data)
        result_df = result_df.merge(g2t_df, left_on = 'mgene_id', right_on = 'gene_id', suffixes=("", "_ri"))
        result_df.drop(columns=["gene_id_ri"],  inplace=True)
        result_df.rename(columns={'transcript_id': "mrna_id"}, inplace=True)
        result_df.to_csv(output, sep='\t', index=False)
    else:
        bind_obj.logger.debug('cis_targets.txt 文件为空, %s' % len(data))
        with open(output, 'w') as out_handler:
            out_handler.write('lncrna_id\tgene_id\tmrna_id\tmgene_id\n')
    bind_obj.logger.debug(' :: export cis_targets.txt file :: completed')

    return output

def export_cis_matrix(data, option_name, dir_path, bind_obj=None):
    """

    :param data: {'task_id': 'lnc_rna', 'exp_level': 'T', 'sg_diff': '5c9da1ba17b2bf206a55d350'}
    :param option_name:
    :param dir_path:
    :param bind_obj:
    :return:
    """
    bind_obj.logger.debug(' :: export cis_targets.txt file')

    conn = db['sg_target_cis_detail']
    recoreds = conn.find({'target_cis_id': ObjectId(data)},
                         {'lncrna_id': 1, 'gene_id': 1, 'mrna_id': 1, 'mgene_id': 1, '_id': 0})
    data = [i for i in recoreds]
    output = os.path.join(dir_path, 'cis_targets.txt')
    if data:
        result_df = pd.DataFrame(data)

        result_df.to_csv(output, sep='\t', index=False)
    else:
        bind_obj.logger.debug('cis_targets.txt 文件为空, %s' % len(data))
        with open(output, 'w') as out_handler:
            out_handler.write('lncrna_id\tgene_id\tmrna_id\tmgene_id\n')
    bind_obj.logger.debug(' :: export cis_targets.txt file :: completed')

    return output

def export_trans_matrix2(data, option_name, dir_path, bind_obj=None):
    """

    :param data: {'task_id': 'lnc_rna', 'exp_level': 'T', 'sg_diff': '5c9da1ba17b2bf206a55d350'}
    :param option_name:
    :param dir_path:
    :param bind_obj:
    :return:
    """
    bind_obj.logger.debug(' :: export trans_targets.txt file')

    # 获取基因转录本对应关系， 只要一组

    task_id = bind_obj.sheet.option("task_id")
    annot_table = db['sg_annotation_query']
    annot_main = annot_table.find_one({"task_id": task_id, "type": "origin", "status": "end"})
    annot_detail = db['sg_annotation_query_detail']
    result = annot_detail.find({"query_id": annot_main["main_id"], "is_gene": True},
                               {"_id": 0, "transcript_id": 1, "gene_id": 1})
    g2t_df = pd.DataFrame(list(result))

    conn = db['sg_target_cistrans_detail']
    recoreds = conn.find({'target_cistrans_id': ObjectId(data), "target_type": "trans"},
                         {'lncrna_id': 1, 'gene_id': 1, 'mgene_id': 1, '_id': 0})

    data = []
    check = defaultdict(int)
    for recored in recoreds:
        lnc_id = recored['lncrna_id']
        num = check[lnc_id]
        if num >= 3:
            continue
        check[lnc_id] += 1
        data.append(recored)

    output = os.path.join(dir_path, 'trans_targets.txt')
    if data:
        result_df = pd.DataFrame(data)
        result_df = result_df.merge(g2t_df, left_on = 'mgene_id', right_on = 'gene_id', suffixes=("", "_ri"))
        result_df.drop(columns=["gene_id_ri"],  inplace=True)
        result_df.rename(columns={'transcript_id': "mrna_id"}, inplace=True)

        result_df.to_csv(output, sep='\t', index=False)
    else:
        bind_obj.logger.debug('trans_targets.txt 文件为空, %s' % len(data))
        with open(output, 'w') as out_handler:
            out_handler.write('lncrna_id\tgene_id\tmrna_id\tmgene_id\n')
    bind_obj.logger.debug(' :: export trans_targets.txt file :: completed')

    return output

def export_trans_matrix(data, option_name, dir_path, bind_obj=None):
    """

    :param data: {'task_id': 'lnc_rna', 'exp_level': 'T', 'sg_diff': '5c9da1ba17b2bf206a55d350'}
    :param option_name:
    :param dir_path:
    :param bind_obj:
    :return:
    """
    bind_obj.logger.debug(' :: export trans_targets.txt file')

    conn = db['sg_target_trans_detail']
    recoreds = conn.find({'target_trans_id': ObjectId(data)},
                         {'lncrna_id': 1, 'gene_id': 1, 'mrna_id': 1, 'mgene_id': 1, '_id': 0})
    data = []
    check = defaultdict(int)
    for recored in recoreds:
        lnc_id = recored['lncrna_id']
        num = check[lnc_id]
        if num >= 3:
            continue
        check[lnc_id] += 1
        data.append(recored)

    output = os.path.join(dir_path, 'trans_targets.txt')
    if data:
        result_df = pd.DataFrame(data)
        result_df.to_csv(output, sep='\t', index=False)
    else:
        bind_obj.logger.debug('trans_targets.txt 文件为空, %s' % len(data))
        with open(output, 'w') as out_handler:
            out_handler.write('lncrna_id\tgene_id\tmrna_id\tmgene_id\n')
    bind_obj.logger.debug(' :: export trans_targets.txt file :: completed')

    return output

def export_multi_gene_list(data, option_name, dir_path, bind_obj=None):
    data = json.loads(data)
    geneset_id_list = data['geneset_id'].split(",")
    source = data['source']
    multi_geneset_path = dir_path + "/multi_geneset_list"
    main_collection = db['sg_geneset']
    collection = db['sg_geneset_detail']
    with open(multi_geneset_path, "wb") as out_handler:
        if len(geneset_id_list) == 1 and source == 'diff_exp':
            geneset_id = geneset_id_list[0]
            my_result = main_collection.find_one({'main_id': ObjectId(geneset_id)})
            if not my_result:
                bind_obj.set_error("意外错误，geneset_id:{}在sg_geneset中未找到！".format(ObjectId(geneset_id)))
            geneset_name = my_result["name"]
            results = collection.find_one({"geneset_id": ObjectId(geneset_id)})
            seq_list = results["seq_list"]
            if 'regulate_list' not in results:
                raise bind_obj.set_error(
                    'geneset_detail table must contain "regulate_list" field when provide one geneset id')
            regulate_list = results["regulate_list"]
            up_list, down_list = [], []
            for seq_id, regulate in zip(seq_list, regulate_list):
                if regulate == 'up':
                    up_list.append(seq_id)
                elif regulate == 'down':
                    down_list.append(seq_id)
                else:
                    continue
            out_handler.write(geneset_name + '_up\t' + ','.join(up_list) + '\n')
            out_handler.write(geneset_name + '_down\t' + ','.join(down_list) + '\n')
        else:
            for n, gi in enumerate(geneset_id_list):
                my_result = main_collection.find_one({'main_id': ObjectId(gi)})
                if not my_result:
                    bind_obj.set_error("意外错误，geneset_id:{}在sg_geneset中未找到！".format(ObjectId(gi)))
                out_handler.write(my_result["name"] + "\t")
                results = collection.find_one({"geneset_id": ObjectId(gi)})
                out_handler.write(",".join(results["seq_list"]) + "\n")

    return multi_geneset_path

def export_go_detail(data, option_name, dir_path, bind_obj=None):
    geneset_id_list = data.split(",")
    multi_geneset_path = dir_path + "/multi_geneset_list"
    main_collection = db['sg_geneset']
    collection = db['sg_geneset_detail']
    with open(multi_geneset_path, "wb") as out_handler:
        if len(geneset_id_list) == 2:
            for n, gi in enumerate(geneset_id_list):
                my_result = main_collection.find_one({'main_id': ObjectId(gi)})
                if not my_result:
                    bind_obj.set_error("意外错误，geneset_id:{}在sg_geneset中未找到！".format(ObjectId(gi)))
                out_handler.write(my_result["name"] + "\t")
                results = collection.find_one({"geneset_id": ObjectId(gi)})
                out_handler.write(",".join(results["seq_list"]) + "\n")
        else:
            geneset_id = geneset_id_list[0]
            my_result = main_collection.find_one({'main_id': ObjectId(geneset_id)})
            if not my_result:
                bind_obj.set_error("意外错误，geneset_id:{}在sg_geneset中未找到！".format(ObjectId(geneset_id)))
            geneset_name = my_result["name"]
            results = collection.find_one({"geneset_id": ObjectId(geneset_id)})
            seq_list = results["seq_list"]
            if 'regulate_list' not in results:
                raise bind_obj.set_error(
                    'geneset_detail table must contain "regulate_list" field when provide one geneset id')
            regulate_list = results["regulate_list"]
            up_list, down_list = [], []
            for seq_id, regulate in zip(seq_list, regulate_list):
                if regulate == 'up':
                    up_list.append(seq_id)
                else:
                    down_list.append(seq_id)
            out_handler.write(geneset_name + '_up\t' + ','.join(up_list) + '\n')
            out_handler.write(geneset_name + '_down\t' + ','.join(down_list) + '\n')

    return multi_geneset_path

def export_kegg_detail(data, option_name, dir_path, bind_obj=None):
    """

    :param data: {''}
    :param option_name:
    :param dir_path:
    :param bind_obj:
    :return:
    """
    geneset_id_list = data.split(",")
    multi_geneset_path = dir_path + "/multi_geneset_list"
    main_collection = db['sg_geneset']
    collection = db['sg_geneset_detail']
    with open(multi_geneset_path, "wb") as out_handler:
        if len(geneset_id_list) == 2:
            for n, gi in enumerate(geneset_id_list):
                my_result = main_collection.find_one({'main_id': ObjectId(gi)})
                if not my_result:
                    bind_obj.set_error("意外错误，geneset_id:{}在sg_geneset中未找到！".format(ObjectId(gi)))
                out_handler.write(my_result["name"] + "\t")
                results = collection.find_one({"geneset_id": ObjectId(gi)})
                out_handler.write(",".join(results["seq_list"]) + "\n")
        else:
            geneset_id = geneset_id_list[0]
            my_result = main_collection.find_one({'main_id': ObjectId(geneset_id)})
            if not my_result:
                bind_obj.set_error("意外错误，geneset_id:{}在sg_geneset中未找到！".format(ObjectId(geneset_id)))
            geneset_name = my_result["name"]
            results = collection.find_one({"geneset_id": ObjectId(geneset_id)})
            seq_list = results["seq_list"]
            if 'regulate_list' not in results:
                bind_obj.set_error(
                    'geneset_detail table must contain "regulate_list" field when provide one geneset id')
            regulate_list = results["regulate_list"]
            up_list, down_list = [], []
            for seq_id, regulate in zip(seq_list, regulate_list):
                if regulate == 'up':
                    up_list.append(seq_id)
                else:
                    down_list.append(seq_id)
            out_handler.write(geneset_name + '_up\t' + ','.join(up_list) + '\n')
            out_handler.write(geneset_name + '_down\t' + ','.join(down_list) + '\n')

    return multi_geneset_path

def export_enrich_heatmap(data, option_name, dir_path, bind_obj=None):
    """
        ['task_id', 'submit_location', 'geneset_type', 'enrich_type',
        'geneset_enrich_ids', 'go_level', 'task_type', 'pvalue_type', 'top_num', 'ids_list', 'go_level']

    :param data:
    :param option_name:
    :param dir_path:
    :param bind_obj:
    :return:
    """
    ids_filer_list = None
    if 'ids_list' in data:
        ids_filer_list = {i for i in data['ids_list']}
    query_dict = {}
    filter_dict = dict(_id=0)
    out_file = os.path.join(dir_path, 'heatmap_data.xls')

    if data['enrich_type'].lower() == 'kegg':
        if data['pvalue_type'].lower() == 'pvalue':
            # if ids_filer_list is None:
            #     query_dict['pvalue'] = {'$lte': data['p_value']}
            filter_dict['pvalue'] = 1
            p_field = 'pvalue'
        else:
            # if ids_filer_list is None:
            #     query_dict['corrected_pvalue'] = {'$lte': data['p_value']}
            filter_dict['corrected_pvalue'] = 1
            p_field = 'corrected_pvalue'

        main_collection = db['sg_geneset_kegg_enrich']
        collection = db['sg_geneset_kegg_enrich_detail']
        ids = [ObjectId(i) for i in data['geneset_enrich_ids'].strip().split(',')]
        filter_dict['id'] = 1
        filter_dict['kegg_type'] = 1
        filter_dict['discription'] = 1

        data_list = {}
        detail_info = defaultdict(dict)
        needed_ids = set()
        geneset_names= []
        min_non_zero = 1
        has_zero = False
        for id_obj in ids:
            query_dict['kegg_enrich_id'] = id_obj
            name = main_collection.find_one({'_id': id_obj})['geneset_name']
            geneset_names.append(name)
            sub_data = {}
            data_list[name] = sub_data
            for dic in collection.find(query_dict, filter_dict):
                kegg_id = dic['id']
                sub_dic = detail_info[kegg_id]
                if 'kegg_type' not in sub_dic:
                    sub_dic['kegg_type'] = dic.pop('kegg_type')
                    sub_dic['discription'] = dic.pop('discription')
                p_value = dic[p_field]
                sub_data[kegg_id] = p_value
                if p_value == 0:
                    has_zero = True
                else:
                    min_non_zero = min(p_value, min_non_zero)
            if ids_filer_list is None:
                # 过滤top num 的ids
                top_num_ids = {t for t in
                               sorted(((k, v) for k, v in sub_data.items()), key=lambda items: items[1])[
                               : data['top_num']]}
                # 过滤p_value值
                p_filter_ids = {t[0] for t in top_num_ids if t[1] <= data['p_value']}
                # 取所有ids的并集
                needed_ids.update(p_filter_ids)

        df = pd.DataFrame(data_list)
        ids_filer_list = ids_filer_list or needed_ids
        if len(ids_filer_list) == 0:
            bind_obj.set_error('提供的id集合和基因集没有交集，请重新输入参数')
        if len(geneset_names) != len(set(geneset_names)):
            bind_obj.set_error('提供的id集合名称重复：' + str(geneset_names))

        df = df.loc[df.index & ids_filer_list]
        if has_zero:  # 0 值取最小正实数的十分之一 [比最小正实数小一个数量级]
            df[df == 0] = min_non_zero / 10
        df.fillna(1, inplace=True)
        df = -np.log10(df)
        df[df == 0] = 0
        indexs = df.index
        df['kegg_type'] = [detail_info[i]['kegg_type'] for i in indexs]
        df['discription'] = [detail_info[i]['discription'] for i in indexs]

        df.index.name = 'kegg_id'
        df.to_csv(out_file, sep='\t', index=True, header=True)
    else:
        if data['pvalue_type'].lower() == 'pvalue':
            # if ids_filer_list is None:
            #     query_dict['p_uncorrected'] = {'$lte': data['p_value']}
            filter_dict['p_uncorrected'] = 1
            p_field = 'p_uncorrected'
        else:
            # if ids_filer_list is None:
            #     query_dict['p_corrected'] = {'$lte': data['p_value']}
            filter_dict['p_corrected'] = 1
            p_field = 'p_corrected'

        main_collection = db['sg_geneset_go_enrich']
        collection = db['sg_geneset_go_enrich_detail']
        ids = [ObjectId(i) for i in data['geneset_enrich_ids'].strip().split(',')]
        filter_dict['go_id'] = 1
        filter_dict['go_type'] = 1
        filter_dict['discription'] = 1
        if 'go_level' in data:
            go_level = data['go_level']
            if go_level.lower() != 'all':
                query_dict['go_type'] = go_level.upper()

        data_list = {}
        detail_info = defaultdict(dict)
        needed_ids = set()
        geneset_names= []
        min_non_zero = 1
        has_zero = False
        for id_obj in ids:
            query_dict['go_enrich_id'] = id_obj
            name = main_collection.find_one({'_id': id_obj})['geneset_name']
            geneset_names.append(name)
            sub_data = {}
            data_list[name] = sub_data
            print(query_dict, '========================')
            for dic in collection.find(query_dict, filter_dict):
                go_id = dic['go_id']
                sub_dic = detail_info[go_id]
                if 'go_type' not in sub_dic:
                    sub_dic['go_type'] = dic.pop('go_type')
                    sub_dic['discription'] = dic.pop('discription')
                p_value = dic[p_field]
                sub_data[go_id] = p_value
                if p_value == 0:
                    has_zero = True
                else:
                    min_non_zero = min(p_value, min_non_zero)

            if ids_filer_list is None:
                # 过滤top num 的ids
                top_num_ids = {t for t in
                               sorted(((k, v) for k, v in sub_data.items()), key=lambda items: items[1])[
                               : data['top_num']]}
                # 过滤p_value值
                p_filter_ids = {t[0] for t in top_num_ids if t[1] <= data['p_value']}
                # 取所有ids的并集
                needed_ids.update(p_filter_ids)

        df = pd.DataFrame(data_list)
        ids_filer_list = ids_filer_list or needed_ids
        if len(ids_filer_list) == 0:
            bind_obj.set_error('提供的id集合和基因集没有交集，请重新输入参数')
        df = df.loc[df.index & ids_filer_list]
        if has_zero:  # 0 值取最小正实数的十分之一
            df[df == 0] = min_non_zero / 10
        df.fillna(1, inplace=True)
        df = -np.log10(df)
        df[df == 0] = 0
        indexs = df.index
        df['go_type'] = [detail_info[i]['go_type'] for i in indexs]
        df['discription'] = [detail_info[i]['discription'] for i in indexs]
        df.fillna(1, inplace=True)

        df.index.name = 'go_id'
        df.to_csv(out_file, sep='\t', index=True, header=True)

    return out_file


import numpy as np
import pandas as pd
from sklearn.svm import SVR
from sklearn.model_selection import GridSearchCV, KFold
from sklearn.metrics import mean_absolute_error, mean_squared_error
import lightgbm as lgb
from utils.load import load_feature_target_set
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
plt.style.use('seaborn')

###############################
###### HELPER FUNCTIONS #######
###############################

def evaluation_mse_score(mod, X, y_true, plot=False, path="figures/", name="test_figure"):
    y_pred = mod.predict(X)
    if plot:
        fig, ax = plt.subplots(1,1, figsize=(10,10))
        ax.scatter(y_pred, y_true)
        ax.legend()
        plt.savefig(path + name + ".png")
    return -mean_squared_error(y_true, y_pred)

def evaluation_mae(mod, X, y_true, plot=False, path="figures/", name="test_figure"):
    y_pred = mod.predict(X)
    if plot:
        fig, ax = plt.subplots(1,1, figsize=(10,10))
        ax.scatter(y_pred, y_true)
        ax.legend()
        plt.savefig(path + name + ".png")
    return mean_absolute_error(y_true, y_pred)

def evaluation_model_torch(mod: nn.Module, data, criterion):
    pred = mod(data.x_t)
    loss = criterion(pred, data.y_t)
    return loss.item()

def stratified_split_by_id(X_recurrent, X_simple, X_baseline, y, ids: list):
    """ the ids list containes the id of each user for each feature/target pair. Must be same length as features/targets """
    from sklearn.model_selection import train_test_split, StratifiedShuffleSplit
    train_idx, test_idx = next(StratifiedShuffleSplit(n_splits=1, train_size=.8, random_state=6).split(ids, ids))

    X_simple, X_recurrent, X_baseline, y = np.asarray(X_simple), np.asarray(X_recurrent), np.asarray(X_baseline), np.array(y)
    X_simple_train, X_simple_test = X_simple[train_idx], X_simple[test_idx]
    X_recurrent_train, X_recurrent_test = X_recurrent[train_idx], X_recurrent[test_idx]
    X_baseline_train, X_baseline_test = X_baseline[train_idx], X_baseline[test_idx]
    y_train, y_test = y[train_idx], y[test_idx]
    return X_simple_train, X_recurrent_train, X_baseline_train, y_train, X_simple_test, X_recurrent_test, X_baseline_test, y_test

########################
###### MODEL RNN #######
########################
from utils.torch_utils import train_torch_model, cross_validate_torch, LSTM, evaluation_model_torch_numpy_mse, evaluation_model_torch_numpy_mae, DataSet

def fit_rnn(args):
    X_recurrent, X_simple, X_baseline, y, ids = load_feature_target_set()
    _, X_train, _, y_train, _, X_test, _, y_test = stratified_split_by_id(X_recurrent, X_simple, X_baseline, y, ids)
    n_features = X_train[0].shape[1]

    # turn off once correct specification determined
    gridsearch = True

    if gridsearch == True:
        param_names = ["lr", "batch_size", "hid_dim", "epochs"]
        columns = param_names + ["avg_score"]
        gridsearch_results = pd.DataFrame(columns = columns)
        k = 0
        for lr in [0.01]:
            for batch_size in [10]:
                for hid_dim in [5, 15, 25, 50]:
                    for epochs in [5,10,20,40]:
                        params = {
                            "lr": lr,
                            "batch_size": batch_size,
                            "hid_dim": hid_dim,
                            "epochs": epochs
                        }
                        fold_scores_train, fold_scores_test = cross_validate_torch(
                            X_train, y_train,
                            verbose = 0,
                            cv = KFold(n_splits=5, shuffle=True, random_state=52),
                            **params
                        )
                        print(params, fold_scores_test)
                        gridsearch_results.loc[k] = [lr, batch_size, hid_dim, epochs, fold_scores_test]
                        k+=1
        lr, batch_size, hid_dim, epochs, _ = tuple(gridsearch_results.sort_values("avg_score").reset_index(drop=True).iloc[0].values)
        batch_size, epochs, hid_dim = int(batch_size), int(epochs), int(hid_dim)

        print(f"RNN - GridSearch - Best params = {dict(zip(columns, (lr, batch_size, hid_dim, epochs)))} - Best avg sore = {_}")
    else:
        # default parameters
        lr = .01
        batch_size = 10
        hid_dim = 5
        epochs = 40

    data_train = DataSet(X_train, y_train)
    data_test = DataSet(X_test, y_test)

    # after cross validation or using default params, train model on full training set and evaluate on the unseen test set
    model = LSTM(in_dim=n_features, hid_dim=hid_dim, out_dim=1)
    model, train_losses, test_losses = train_torch_model(
        model = model,
        optimizer = torch.optim.Adam(model.parameters(), lr=lr),
        criterion = nn.MSELoss(),
        data_train = data_train,
        data_validate = data_test,
        epochs = epochs
    )
    best_loss_test = test_losses[-1]

    mae = evaluation_model_torch_numpy_mae(model, data_test)

    y_pred = model(data_test.x_t).detach().numpy().reshape(-1,)
    y_true = data_test.y_t.detach().numpy().reshape(-1,)

    print(f"RNN - Test - MSE loss = {best_loss_test} - MAE loss = {mae}")

    return y_pred, y_true

########################
###### MODEL LGB #######
########################

def fit_lgb(args):
    """ Determine the model specification through a cross validated grid search"""
    X_recurrent, X_simple, X_baseline, y, ids = load_feature_target_set()
    columns = X_simple.columns
    X_train, _, _, y_train, X_test, _, _, y_test = stratified_split_by_id(X_recurrent, np.array(X_simple), X_baseline, y, ids)

    gridsearch = False
    if gridsearch:
        param_grid = {
            # "num_leaves": [],
            "learning_rate": [0.16, 0.08, 0.04, 0.02, 0.01],
            "max_depth": [-1, 1, 2, 4],
            "num_leaves": [5, 10, 20, 40, 80],
            "n_estimators": [50, 100],
            "min_split_gain": [0, 1]
        }
    else:
        param_grid = {
            # "num_leaves": [],
            "learning_rate": [0.02,],
            "max_depth": [-1,],
            "num_leaves": [20,],
            "n_estimators": [100]
        }
    mod = lgb.LGBMModel(objective="regression")
    gs = GridSearchCV(
        estimator = mod,
        param_grid = param_grid,
        cv = KFold(5, shuffle=True, random_state=52),
        scoring = evaluation_mse_score,
        verbose=0
    )
    gs.fit(pd.DataFrame(columns = columns, data = X_train), y_train)
    y_pred = gs.predict(X_test)

    # REMOVE THIS
    lgb.plot_importance(gs.best_estimator_, importance_type="gain")
    plt.savefig("figures/lgb_importance.png", dpi = 250, bbox_inches = "tight")
    plt.clf()
    # REMOVE THIS

    mse = mean_squared_error(y_test, y_pred)
    mae = mean_absolute_error(y_test, y_pred)
    print(f"LGBM - GridSearch - Best params = {gs.best_params_} - Best avg score = {gs.best_score_}")
    print(f"LGBM - Test - MSE loss = {mse} - MAE loss = {mae}")

    return y_pred, y_test

#############################
###### MODEL BASELINE #######
#############################
def fit_baseline(args):
    """ the baseline model is an AR(1) with lagged coefficient = 1 """
    X_recurrent, X_simple, X_baseline, y, ids = load_feature_target_set()
    _, _, X_train, y_train, _, _, X_test, y_test = stratified_split_by_id(X_recurrent, X_simple, X_baseline, y, ids)
    y_pred = X_test.copy()

    mse = mean_squared_error(y_test, y_pred)
    mae = mean_absolute_error(y_test, y_pred)
    print(f"Baseline - Test - MSE loss = {mse} - MAE loss = {mae}")

    return y_pred, y_test

##########################
###### MODEL ARIMA #######
##########################
def fit_arima(args):
    """ the baseline model is an AR(1) with coefficient = 1 """
    X_recurrent, X_simple, X_baseline, y, ids = load_feature_target_set()
    _, X_train, _, y_train, _, X_test, _, y_test = stratified_split_by_id(X_recurrent, X_simple, X_baseline, y, ids)

    print(X_train.shape)
    X_recurrent = np.array([x["mean_mood"] for x in X_train]).reshape(len(X_train), len(X_train[0]))
    # TODO add ols
    import statsmodels.api as sm
    mod = sm.OLS(y_train, X_train)
    y_pred = mod.predict(X_test)

    mse = mean_squared_error(y_test, y_pred)
    mae = mean_absolute_error(y_test, y_pred)
    print(f"Baseline - Test - MSE loss = {mse} - MAE loss = {mae}")

#######################
## GENERAL FUNCTIONS ##
#######################
def evaluate_model(model, X_test, y_test):
    y_pred = model.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)
    mae = mean_absolute_error(y_test, y_pred)
    return mse, mae

model_functions = {
    "baseline": fit_baseline,
    "lgb": fit_lgb,
    "rnn": fit_rnn,
    "arima": fit_arima
}


from datetime import datetime
import collections
import struct
import serial
import six.moves
import logging
import math

import modbus_tk.defines as modbus_constants
import modbus_tk.modbus_tcp as modbus_tcp
import modbus_tk.modbus_rtu as modbus_rtu
from modbus_tk.exceptions import ModbusError

from . import helpers

logger = logging.getLogger(__name__)

# In cache representation of modbus field.
Datum = collections.namedtuple('Datum', ('value', 'timestamp'))

class ModbusFieldException(Exception):
    pass

class Field:
    """Describes/defines a logical modbus field.

        A logical field may be mapped to one or more modbus registers or coils.

        Field implements the 'descriptor' pattern with __get__ and __set__ methods.

        A field defines the modbus table, address offset, length and datatype of the
        logical field.  Getting and setting of field values is as simple as
        attribute get/set (via to descriptor pattern).

        myModbusClientInstance.float_field_1 = 3.14159
        ...
        myModbusClientInstance.write_all()    # Writes all modified fields to the modbus slave.

        Getting the value is simpler:

        myModbusClientInstance.float_field_1  #Reads from cache or from modbus slave.

    """

    def __init__(self, name, address, datatype, units, precision, transform, table, op_mode, mixed=False):
        self._name = name
        self._address = address
        self._type = datatype
        self._units = units
        self._precision = precision
        self._transform = transform
        self._table = table
        self._op_mode = op_mode
        self._mixed = mixed

    @property
    def is_struct_format(self):
        """Returns True if the type of this field is described by a struct
            format string, eg:  ">h", instead of one of the field tuples.
        """
        return isinstance(self._type, str) and len(self._type) and \
            self._type[0] in (helpers.BIG_ENDIAN, helpers.LITTLE_ENDIAN)

    @property
    def format_string(self):
        """Generate a struct format string from the type, which may be a tuple
           like:  ("I", 2, 1)   # representing UINT
           or an actual struct format strng like:  ">2h"
        """
        if isinstance(self._type, tuple):
            try:
                return str(self.length) + self._type[helpers.FORMAT]
            except TypeError:
                logger.warning("Undefined type for {0}".format(self))
        elif isinstance(self._type, str):
            return self._type[1:]
        else:
            raise ModbusFieldException("Unknown type string for {0}".format(self))

    @property
    def byte_order(self):
        """Byte order is None (default) unless self._type is a format string
        with the first character specifying byte order.

        :return: byte_order is specified or None
        """
        if self.is_struct_format:
            return self._type[0]

        return None

    @property
    def length(self):
        if isinstance(self._type, tuple):
            return self._type[helpers.LENGTH]
        # Returning length of 1 on struct types (non tuples) which means that they
        # will not be converted into list by modbus client
        return 1

    @property
    def table(self):
        return self._table

    @property
    def address(self):
        return self._address

    @property
    def name(self):
        return self._name

    @property
    def type(self):
        return self._type

    @property
    def units(self):
        return self._units

    @property
    def precision(self):
        return self._precision

    @property
    def transform(self):
        return self._transform

    @property
    def is_array_field(self):
        return self.length > 1 and self._type[helpers.FORMAT] != 's'

    @property
    def mixed(self):
        return self._mixed

    def value_for_transport(self, value, modbus_client=None):
        """
            Modbus deals only in 2-byte registers.  4-byte types must be
            sent as 2 x 2-byte values. These are created using the
            Python struct library.

            Before any unpacking, we transform the value if its transform
            function has defined an inverse.

            :param value: the value to be transformed into a binary string

        """
        transformed_value = value
        transform_args = [value]
        if hasattr(self._transform, 'inverse'):
            try:
                if hasattr(self._transform, 'register_args'):
                    for reg_name in self._transform.register_args:
                        transform_args.append(getattr(modbus_client, reg_name))

                transformed_value = self._transform.inverse(*transform_args)
                if not isinstance(transformed_value, type(value)):
                    # Inverse causes integer values to be converted to float,
                    # convert value to original datatype and pass it back
                    transformed_value = type(value)(transformed_value)
            except ZeroDivisionError:
                transformed_value = 0
        return transformed_value

    @classmethod
    def default_holding_register(cls,  name, address, type, units, transform):
        """Shortcut for common definition format."""
        return cls(name, address, type, units, 1, transform, helpers.REGISTER_READ_WRITE, helpers.OP_MODE_READ_WRITE)

    @property
    def absolute_address(self):
        return self.address + helpers.TABLE_ADDRESS[self._table]

    def transform_value(self, value, modbus_client=None):
        """
            Return the transformed value of the field or just
            the value if no transform has been defined.

        :param value: value to be transformed.
        :return:
        """
        transform_args = [value]
        if value == 0:
            return value
        if self._transform:
            if hasattr(self._transform, 'register_args'):
                for reg_name in self._transform.register_args:
                    transform_args.append(getattr(modbus_client, reg_name))

            return self._transform(*transform_args)
        return None

    @property
    def writable(self):
        return self._table in (helpers.COIL_READ_WRITE, helpers.REGISTER_READ_WRITE) and \
               self._op_mode in (helpers.OP_MODE_READ_WRITE, helpers.OP_MODE_WRITE_ONLY)

    @property
    def single_write_function_code(self):
        if self._table == helpers.REGISTER_READ_WRITE:
            return modbus_constants.WRITE_SINGLE_REGISTER
        elif self._table == helpers.COIL_READ_WRITE:
            return modbus_constants.WRITE_SINGLE_COIL
        else:
            raise Exception("Unknown table type or write operation not supported : {0}".format(self._table))

    def __get__(self, instance, owner):
        """
        Return fields data value.
        :param instance: Modbus client
        :param owner: Field object
        :return: Field value from modbus client.
        """
        datum = instance.get_data(self)
        if datum is None or (datetime.utcnow() - datum.timestamp).total_seconds()*1000 >= instance.latency:
            instance.fetch_field(self)
            datum = instance.get_data(self)
        if datum:
            value = Field.convert_mixed(self.type, datum.value) if self._mixed else datum.value
            return self.transform_value(value, instance)
        return None

    def __set__(self, instance, value):
        # If value is None, its a No Op, the field is not updated
        if value is not None:
            if not isinstance(value, str) and value < 0 and self._type in (helpers.USHORT, helpers.UINT, helpers.UINT64):
                raise ValueError("Attempting to assign negative value to unsigned type.")
            if not instance._ignore_op_mode and self._op_mode == helpers.OP_MODE_READ_ONLY:
                raise ValueError("Attempting to write read-only field.")
            value = self.value_for_transport(value, instance)
            if self._mixed:
                value = Field.convert_mixed(self.type, value)
            instance._pending_writes[self] = value

    @staticmethod
    def convert_mixed(datatype, value):
        """Reverse order of register

        :param datatype: register type
        :param value: register value to reverse
        """
        try:
            datatype = datatype[1:] if datatype.startswith((">", "<")) else datatype
            parse_struct = struct.Struct(">{}".format(datatype))
        except AttributeError:
            parse_struct = struct.Struct(">{}".format(datatype[0]))

        value_bytes = parse_struct.pack(value)
        register_values = []
        for i in range(0, len(value_bytes), 2):
            register_values.extend(struct.unpack(">H", value_bytes[i:i + 2]))
        register_values.reverse()
        convert_bytes = bytes.join(b'', [struct.pack(">H", i) for i in register_values])
        return parse_struct.unpack(convert_bytes)[0]

    def fix_address(self, address_style):
        # Translate modbus addressing to absolute offsets
        if address_style in (helpers.ADDRESS_MODBUS, helpers.ADDRESS_OFFSET_PLUS_ONE):

            if address_style == helpers.ADDRESS_MODBUS:
                self._address = self._address - helpers.TABLE_ADDRESS[self._table]
            elif address_style == helpers.ADDRESS_OFFSET_PLUS_ONE:
                self._address = self._address - 1
            if self._address < 0 or self._address > 10000:
                raise Exception("Modbus address out of range for table.")

class Request:
    """
        Represents a contiguous set of logical fields, registers or coils.  The first
        field starts at self.address.

        Request is used to query a collection of fields, possibly of different datatypes,
        in a single modbus request.

    """

    def __init__(self, first_field, data_format='>'):
        self._table = first_field.table
        self._address = first_field.address
        self._next_address = first_field.address
        self._name = first_field.name
        self._count = 0
        self._data_format = first_field.byte_order or data_format
        self._fields = list()

        self.add_field(first_field)

    def table_name(self):
        """Returns display name of modbus table to be accessed by request"""

        if self._table == helpers.REGISTER_READ_WRITE:
            return "REGISTER_READ_WRITE"
        elif self._table == helpers.REGISTER_READ_ONLY:
            return "REGISTER_READ_ONLY"
        elif self._table == helpers.COIL_READ_WRITE:
            return "COIL_READ_WRITE"
        elif self._table == helpers.COIL_READ_ONLY:
            return "COIL_READ_ONLY"
        else:
            return "UNKNOWN"

    def __str__(self):
        return "<Request: {0} - {1}[{2}] x {3} | {4}".format(
            self._name,
            self.table_name(),
            self._address,
            self._count,
            self._data_format
        )

    @property
    def name(self):
        return self._name

    @property
    def fields(self):
        return self._fields

    @property
    def address(self):
        return self._address

    @property
    def count(self):
        return self._count

    @property
    def table(self):
        return self._table

    @property
    def read_function_code(self):
        """Returns a modbus read function code appropriate for the table."""

        if self._table == helpers.REGISTER_READ_WRITE:
            modbus_function_code = modbus_constants.READ_HOLDING_REGISTERS
        elif self._table == helpers.REGISTER_READ_ONLY:
            modbus_function_code = modbus_constants.READ_INPUT_REGISTERS
        elif self._table == helpers.COIL_READ_WRITE:
            modbus_function_code = modbus_constants.READ_COILS
        elif self._table == helpers.COIL_READ_ONLY:
            modbus_function_code = modbus_constants.READ_DISCRETE_INPUTS
        else:
            raise Exception("Unknown table type : {0}".format(self._table))

        return modbus_function_code

    @property
    def write_function_code(self):
        """Returns a modbus write function code appropriate for the table."""

        if self._table == helpers.REGISTER_READ_WRITE:
            # WRITE_MULTIPLE is used unless client cannot support it, see _write_single_values in Client.__init__
            return modbus_constants.WRITE_MULTIPLE_REGISTERS
        elif self._table == helpers.COIL_READ_WRITE:
            # @TODO Temp hack around COIL write problem.
            return modbus_constants.WRITE_SINGLE_COIL if self._count == 1 \
                else modbus_constants.WRITE_MULTIPLE_COILS
        else:
            raise Exception("Unknown table type or write operation not supported : {0}".format(self._table))

    @property
    def formatting(self):
        """
            Provide a struct library format string for non-COIL data types.

            For COILS, we leave the formatting to modbus-tk.
        """
        if self._table not in (helpers.COIL_READ_ONLY, helpers.COIL_READ_WRITE):
            return self._data_format
        else:
            return None

    def able_to_add(self, field):
        return self._table == field.table and \
           self._next_address == field.address and \
           self._count + math.ceil(struct.calcsize(field.format_string) / 2.0) < 124 and \
           field.length == 1 and not field.byte_order and \
           not field.is_struct_format

    def add_field(self, field):
        """Add field to request if it is compatible and contiguous
        otherwise raise

        :return:
        """
        struct_format = field.format_string
        struct_size = struct.calcsize(struct_format)
        if struct_size % 2 == 1:
            struct_size += 1
        self._data_format += struct_format
        # current_request.count += f.type[SIZE] * f.type[LENGTH]
        self._count += struct_size // 2 # We need count to remain a int
        self._fields.append(field)
        self._next_address += struct_size // 2 # We need count to remain a int

    def block_info(self):
        return self._name, self._table, self._address, self._count

    def parse_values(self, results):
        """Parse field values out of the modbus_tk results

        :param results: unpacked results from modbus_tk execute()
        :return: a dictionary of {field_name: value, ...}
        """
        now = datetime.utcnow()
        field_values = {}
        if len(self.fields) == 1 and self.fields[0].is_array_field:
            field = self.fields[0]
            # Array
            field_values = {
                field: Datum([r for r in results], now)
            }
        else:
            # Struct formatted registers and processed as a single field.
            if self.fields and self.fields[0].is_struct_format:
                if type(results) is list or type(results) is tuple:
                    if len(results) > 1:
                        results = (results,)
            # Everything else))
            field_values = collections.OrderedDict(
                [(field, Datum(value, now)) for field, value in six.moves.zip(self.fields, results)]
            )
        return field_values

    @classmethod
    def compile_requests(cls, fields, byte_order):
        """

        Creates a set of Modbus requests for the fields provided.  The fields
        are sorted by table and address so that a minimum number of
        requests can be created.

        These requests are used for both reading and writing.

        :param fields: List of fields sorted by address.
        :param byte_order: Byte order of the modbus slave.
        :return: List of Requests
        """
        requests = list()

        if not fields:
            return requests

        fields.sort(key=lambda f: f.table * 100000 + f.address)

        # Set up the initial request before starting the loop
        current_request = None

        for f in fields:
            # Decide if we need to start a new request

            if current_request is None or not current_request.able_to_add(f):
                current_request = Request(f, data_format=byte_order)
                requests.append(current_request)
                if f.is_struct_format or f.is_array_field:
                    current_request = None
            else:
                current_request.add_field(f)

        return requests

# WARNING: Currently the modbus_tk library is not able to make connections from 2 masters on one host to 2 slaves on
# one host - this will will prevent a single platform from being able to communicate to 2 slaves on one IP as each
# instance of a Modbus_Tk driver creates a new Modbus master.
# Issue on Modbus_Tk Github: https://github.com/ljean/modbus-tk/issues/124
class Client:

    """
    Generic modbus master.  It functions as a traditional "client" making requests of
    a modbus slave which acts like a traditional "server".

    Subclass this class with Field class variables.  See module description above.
    """

    byte_order = helpers.BIG_ENDIAN
    addressing = helpers.ADDRESS_OFFSET

    __meta = None

    @classmethod
    def _build_meta(cls):
        if cls.__meta is None:
            meta = dict()
            for name, obj in cls.__dict__.items():
                if type(obj) == Field:
                    obj.fix_address(getattr(cls, 'addressing'))
                    meta[name] = obj

            # Maintain a list of fields sorted by address (ascending)
            meta[helpers.META_FIELDS] = list(meta.values())                         # Turns Python3 view into a list.
            meta[helpers.META_FIELDS].sort(key=lambda f: f.address)
            meta[helpers.META_REQUESTS] = Request.compile_requests(meta[helpers.META_FIELDS], cls.byte_order)
            # Dictionary for easy lookup of the request that corresponds to a field.
            meta[helpers.META_REQUEST_MAP] = {field: request for request in meta[helpers.META_REQUESTS]
                                              for field in request._fields}
            cls.__meta = meta
        return cls.__meta

    def __init__(self, *args, **kwargs):
        """
            Sets up the Modbus Master to communicate with a slave.  Supports ModbusTCP and ModbusRTU depending
            on the values provided in device_address and port.  If port == None, ModbusRTU is used.

        :param device_address:  IP address if using ModbusTCP, a device id (eg /dev/ttyAMA0) if using ModbusRTU
        :param port: Port number for ModbusTCP, None if using ModbusRTU
        :param slave_address: Unique ID of the Modbus slave.
        :param baud: Baud rate for ModbusRTU.
        :param latency: Controls the cache refresh rate.  Max age of data values read from slave.
        :param ignore_op_mode: Turns read/write error checking off or on.  Deprecated.
        :param timeout_in_sec: Time to wait for a response from the slave.
        :param verbose:
        :param write_single_values: Write registers or coils one value at a time (WRITE_SINGLE_REGISTER, etc.).
        :return:
        """
        # Build up metadata dictionaries from the Fields defined on the class
        self._build_meta()

        # Support for legacy usage where device_address, port are passed in constructor.
        device_address = None if len(args) == 0 else args[0]
        port = None if len(args) < 2 or args[1] is None else int(args[1])
        self.slave_address = None if len(args) < 3 or args[2] is None else int(args[2])

        # Optional keyword arguments
        if self.slave_address is None:
            self.slave_address = kwargs.pop('slave_address', 1)
        self.latency = kwargs.pop('latency', 1000)
        self._ignore_op_mode = kwargs.pop('ignore_op_mode', False)  # Protection against unintended writes

        # Some modbus clients do not support the WRITE_MULTIPLE_REGISTERS function call.
        self._write_single_values = kwargs.pop('write_single_values', False)

        baud = kwargs.pop('baudrate', 19200)
        bytesize = kwargs.pop('bytesize', 8)
        parity = kwargs.pop('parity', 'N')
        stopbits = kwargs.pop('stopbits', 1)
        timeout_in_sec = kwargs.pop('timeout_in_sec', 1.0)
        verbose = kwargs.pop('verbose', False)

        if device_address:
            if port is not None:
                self.client = modbus_tcp.TcpMaster(host=device_address, port=port, timeout_in_sec=timeout_in_sec)
            else:
                self.client = modbus_rtu.RtuMaster(
                    serial.Serial(device_address,
                                  baudrate=baud,
                                  bytesize=bytesize,
                                  parity=parity,
                                  stopbits=stopbits,
                                  xonxoff=0)
                )
                self.client.set_timeout(1.0)
            self.client.set_verbose(verbose)

        self._data = collections.OrderedDict()
        self._pending_writes = dict()
        self._error_count = 0

    def set_transport_tcp(self, hostname, port, timeout_in_sec=1.0):
        self.client = modbus_tcp.TcpMaster(host=hostname, port=int(port), timeout_in_sec=timeout_in_sec)
        return self

    def set_transport_rtu(self, device, baudrate, bytesize, parity, stopbits, xonxoff):
        self.client = modbus_rtu.RtuMaster(
            serial.Serial(device,
                          baudrate=baudrate, bytesize=bytesize, parity=parity, stopbits=stopbits, xonxoff=xonxoff,
                          rtscts=False, writeTimeout=None, dsrdtr=False, interCharTimeout=None)
        )
        self.client.set_timeout(1.0)
        return self

    def __str__(self):
        return 'data: {0}\npending_writes: {1}'.format(self._data, self._pending_writes)

    def pprint(self):
        response = "pending writes: \n"
        for f, v in six.iteritems(self._pending_writes):
            response = response + "\t{0} : {1}".format(f, v)
        return response

    @property
    def has_pending_writes(self):
        return bool(self._pending_writes)

    def requests(self):
        return self.__meta[helpers.META_REQUESTS]

    def fields(self):
        return self.__meta[helpers.META_FIELDS]

    def field_by_name(self, name):
        return self.__meta.get(name, None)

    @property
    def _fields(self):
        return self.__meta['']

    def get_data(self, field):
        return self._data.get(field, None)

    @property
    def write_single_values(self):
        return self._write_single_values

    def get_request(self, field):
        return self.__meta[helpers.META_REQUEST_MAP].get(field, None)

    def read_request(self, request):
        logger.debug("Requesting: %s", request)
        try:
            results = self.client.execute(
                self.slave_address,
                request.read_function_code,
                request.address,
                quantity_of_x=request.count,
                data_format=request.formatting,
                threadsafe=False
            )
            self._data.update(request.parse_values(results))
        except (AttributeError, ModbusError) as err:
            if "Exception code" in err.message:
                raise Exception("{0}: {1}".format(err.message,
                                                  helpers.TABLE_EXCEPTION_CODE.get(err.message[-1], "UNDEFINED")))
            logger.warning("modbus read_all() failure on request: %s\tError: %s", request, err)

    def read_all(self):
        requests = self.__meta[helpers.META_REQUESTS]
        self._data.clear()
        for r in requests:
            self.read_request(r)

    def dump_all(self):
        self.read_all()
        return [(f,
                 f.transform_value(Field.convert_mixed(f.type, d.value), self) if f.mixed else f.transform_value(d.value, self),
                 d.timestamp) for f, d in six.iteritems(self._data)]

    def write_all(self):
        logger.debug("In write_all")
        fields = list(self._pending_writes.keys())
        if self.write_single_values:
            # Convert values if necessary for transport as modbus supported types.
            for f in fields:
                value = self._pending_writes.pop(f)
                values = list()
                if type(value) in (list, tuple):
                    values.extend(value)
                else:
                    values.append(value)
                logger.debug("Writing modbus data for field %s: %s", f.name, values)
                self.client.execute(
                    self.slave_address,
                    f.single_write_function_code,
                    f.address,
                    quantity_of_x=len(values),
                    output_value=values[0] if len(values) == 1 else values,
                    threadsafe=False
                )
        else:
            requests = Request.compile_requests(fields, self.byte_order)
            for r in requests:
                values = list()
                # Convert values if necessary for transport as modbus supported types.
                for f in r.fields:
                    value = self._pending_writes.pop(f)
                    if type(value) in (list, tuple):
                        values.extend(value)
                    else:
                        values.append(value)
                # Temp workaround for COIL write problem.
                values = [value.encode('utf-8') if isinstance(value, str) else value for value in values]
                if r.write_function_code == modbus_constants.WRITE_SINGLE_COIL:
                    values = values[0]
                self.client.execute(
                    self.slave_address,
                    r.write_function_code,
                    r.address,
                    quantity_of_x=r.count,
                    output_value=values,
                    data_format=r.formatting,
                    threadsafe=False
                )
        if self._pending_writes:
            logger.warning("Did not write ALL values!")
        self._pending_writes.clear()
        self._data.clear()

    def fetch_field(self, field):
        """
            Make a modbus request for the block containing field.
        """
        request = self.get_request(field)
        if request:
            self.read_request(request)

    def close(self):
        self.client.close()


#! /usr/bin/env python3
# encoding: utf-8
#
# (C) 2016 Chris Liechti <cliechti@gmx.net>
#
# SPDX-License-Identifier:    BSD-3-Clause
import socket
import sys
import gc
import micropython
from .connection import Connection
import ulog
gc.collect()

log = ulog.Logger('HTTPD: ')

class Server(object):
    def __init__(self, app, port=80):
        self.app = app

        self.listening_socket = socket.socket()

        # Binding to all interfaces - server will be accessible to other hosts!
        ai = socket.getaddrinfo('0.0.0.0', port)
        log.info('Bind address info: {}'.format(ai))
        addr = ai[0][-1]

        try:
            self.listening_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        except AttributeError:
            log.warn('setsockopt not supported')
        self.listening_socket.bind(addr)
        self.listening_socket.listen(1)

    def wait_for_client(self):
        client_socket, client_addr = self.listening_socket.accept()
        #~ client_socket.settimeout(5)
        try:
            connection = Connection(
                '{}:{}'.format(*client_addr).encode('utf-8'),
                client_socket.makefile('rb'),
                client_socket.makefile('wb'),
                #~ keep=True)
                keep=False)
            connection.do_request(self.app)
        finally:
            #~ print("terminate")
            gc.collect()
            client_socket.close()

    def loop(self):
        self.app.optimize_routes()
        gc.collect()
        micropython.mem_info(1)
        while True:
            try:
                self.wait_for_client()
                #~ socket.print_pcbs()
                #~ gc.collect()
            except OSError as e:
                sys.print_exception(e)
                #~ if 'ENOMEM' in str(e):
                    #~ micropython.mem_info(1)
                    #~ socket.print_pcbs()
                    #~ machine.reset()
            except Exception as e:
                #~ raise
                sys.print_exception(e)


from aiohttp import web
import json

from app import mpdservices

# Обработка запроса информации о файлах
async def get_files_info_handler(request):
    directory = request.rel_url.query.get('directory', None)
    data = await mpdservices.get_directory_files(directory)
    return web.json_response(data)

# Запрос на обновление базы данных
async def update_database_handler(request):
    data = await mpdservices.update_mpd_database()
    return web.json_response(data)

# Запрос на добавление песни в один из сохраненных плейлистов
async def song_to_playlist_handler(request):
    data = await request.json()
    await mpdservices.song_to_playlist(data['file'], data['playlist'])
    return web.Response(status=200)

# Запрос на добавление песни в список воспроизведения
async def song_to_current_playlist_handler(request):
    data = await request.json()
    await mpdservices.to_current_playlist(data['file'])
    return web.Response(status=200)

# Запрос на проигрывание выбранной песни
async def play_song_handler(request):
    pass