import os
import csv
import queue
import logging
import argparse
import traceback
import itertools
import numpy as np

import tensorflow.compat.v1 as tf

from fedlearner.trainer.bridge import Bridge
from fedlearner.model.tree.tree import BoostingTreeEnsamble
from fedlearner.trainer.trainer_master_client import LocalTrainerMasterClient
from fedlearner.trainer.trainer_master_client import DataBlockInfo

'''
目前不太理解的地方：worker、verbosity、max-bins、ignore-fields
'''
def create_argument_parser():
    parser = argparse.ArgumentParser(
        description='FedLearner Tree Model Trainer.')
    #训练角色，leader还是follower
    parser.add_argument('role', type=str,
                        help="Role of this trainer in {'local', "
                             "'leader', 'follower'}")
    #监听本地地址，ip+port
    parser.add_argument('--local-addr', type=str,
                        help='Listen address of the local bridge, ' \
                             'in [IP]:[PORT] format')
    #同伴地址，ip+port
    parser.add_argument('--peer-addr', type=str,
                        help='Address of peer\'s bridge, ' \
                             'in [IP]:[PORT] format')
    #分布式训练时，应用程序的id，默认空
    parser.add_argument('--application-id', type=str, default=None,
                        help='application id on distributed ' \
                             'training.')
    #current worker的排名，等级，默认0
    parser.add_argument('--worker-rank', type=int, default=0,
                        help='rank of the current worker')
    #总的worker数量，默认1
    parser.add_argument('--num-workers', type=int, default=1,
                        help='total number of workers')
    #mode，可以为 train,test,eval,默认为train
    parser.add_argument('--mode', type=str, default='train',
                        help='Running mode in train, test or eval.')
    #数据文件的路径
    parser.add_argument('--data-path', type=str, default=None,
                        help='Path to data file.')
    #验证数据文件的路径，仅用于test模式
    parser.add_argument('--validation-data-path', type=str, default=None,
                        help='Path to validation data file. ' \
                             'Only used in train mode.')
    #bool变量，默认为false，预测不需要数据
    parser.add_argument('--no-data', type=bool, default=False,
                        help='Run prediction without data.')
    #使用的文件扩展
    parser.add_argument('--file-ext', type=str, default='.csv',
                        help='File extension to use')
    #输入文件类型
    parser.add_argument('--file-type', type=str, default='csv',
                        help='input file type: csv or tfrecord')
    #加载已存储模型的路径
    parser.add_argument('--load-model-path',
                        type=str,
                        default=None,
                        help='Path load saved models.')
    #存储输出模型的位置
    parser.add_argument('--export-path',
                        type=str,
                        default=None,
                        help='Path to save exported models.')
    #保存模型的检查点
    parser.add_argument('--checkpoint-path',
                        type=str,
                        default=None,
                        help='Path to save model checkpoints.')
    #存储预测输出的路径
    parser.add_argument('--output-path',
                        type=str,
                        default=None,
                        help='Path to save prediction output.')
    #控制打印日志的数量，默认为1
    parser.add_argument('--verbosity',
                        type=int,
                        default=1,
                        help='Controls the amount of logs to print.')
    #损失函数的选择，默认为logistic
    parser.add_argument('--loss-type',
                        default='logistic',
                        choices=['logistic', 'mse'],
                        help='What loss to use for training.')
    #学习率，梯度下降中的步长，默认为0.3
    parser.add_argument('--learning-rate',
                        type=float,
                        default=0.3,
                        help='Learning rate (shrinkage).')
    #boost 迭代次数，默认为5
    parser.add_argument('--max-iters',
                        type=int,
                        default=5,
                        help='Number of boosting iterations.')
    #决策树的最大深度，默认为3
    parser.add_argument('--max-depth',
                        type=int,
                        default=3,
                        help='Max depth of decision trees.')
    #L2正则化参数，默认为1.0
    parser.add_argument('--l2-regularization',
                        type=float,
                        default=1.0,
                        help='L2 regularization parameter.')
    #最大的直方图维度
    parser.add_argument('--max-bins',
                        type=int,
                        default=33,
                        help='Max number of histogram bins.')
    #并行线程的数量，默认1
    parser.add_argument('--num-parallel',
                        type=int,
                        default=1,
                        help='Number of parallel threads.')
    #bool类型，如果被设置为true，数据第一列被认为是双方都匹配的example id
    parser.add_argument('--verify-example-ids',
                        type=bool,
                        default=False,
                        help='If set to true, the first column of the '
                             'data will be treated as example ids that '
                             'must match between leader and follower')
    #通过名字来忽略数据域，默认空字符串
    parser.add_argument('--ignore-fields',
                        type=str,
                        default='',
                        help='Ignore data fields by name')
    #分类特征的字段名称，特征的值应该为非负整数
    parser.add_argument('--cat-fields',
                        type=str,
                        default='',
                        help='Field names of categorical features. Feature'
                             ' values should be non-negtive integers')
    #是否使用流传输，默认为否
    parser.add_argument('--use-streaming',
                        type=bool,
                        default=False,
                        help='Whether to use streaming transmit.')
    #是否发送预测评分给follower，默认为false
    parser.add_argument('--send-scores-to-follower',
                        type=bool,
                        default=False,
                        help='Whether to send prediction scores to follower.')
    #是否发送指标（metrics）给follower，默认为follower
    parser.add_argument('--send-metrics-to-follower',
                        type=bool,
                        default=False,
                        help='Whether to send metrics to follower.')

    return parser


def parse_tfrecord(record):
    example = tf.train.Example()
    example.ParseFromString(record)

    parsed = {}
    for key, value in example.features.feature.items():
        kind = value.WhichOneof('kind')
        if kind == 'float_list':
            assert len(value.float_list.value) == 1, "Invalid tfrecord format"
            parsed[key] = value.float_list.value[0]
        elif kind == 'int64_list':
            assert len(value.int64_list.value) == 1, "Invalid tfrecord format"
            parsed[key] = value.int64_list.value[0]
        elif kind == 'bytes_list':
            assert len(value.bytes_list.value) == 1, "Invalid tfrecord format"
            parsed[key] = value.bytes_list.value[0]
        else:
            raise ValueError("Invalid tfrecord format")

    return parsed


def extract_field(field_names, field_name, required):
    if field_name in field_names:
        return []

    assert not required, \
        "Field %s is required but missing in data"%field_name
    return None


def read_data(file_type, filename, require_example_ids,
              require_labels, ignore_fields, cat_fields):
    logging.debug('Reading data file from %s', filename)

    if file_type == 'tfrecord':
        reader = tf.io.tf_record_iterator(filename)
        reader, tmp_reader = itertools.tee(reader)
        first_line = parse_tfrecord(next(tmp_reader))
        field_names = first_line.keys()
    else:
        fin = tf.io.gfile.GFile(filename, 'r')
        reader = csv.DictReader(fin)
        field_names = reader.fieldnames

    example_ids = extract_field(
        field_names, 'example_id', require_example_ids)
    raw_ids = extract_field(
        field_names, 'raw_id', False)
    labels = extract_field(
        field_names, 'label', require_labels)

    ignore_fields = set(filter(bool, ignore_fields.strip().split(',')))
    ignore_fields.update(['example_id', 'raw_id', 'label'])
    cat_fields = set(filter(bool, cat_fields.strip().split(',')))
    for name in cat_fields:
        assert name in field_names, "cat_field %s missing"%name

    cont_columns = list(filter(
        lambda x: x not in ignore_fields and x not in cat_fields, field_names))
    cont_columns.sort(key=lambda x: x[1])
    cat_columns = list(filter(
        lambda x: x in cat_fields and x not in cat_fields, field_names))
    cat_columns.sort(key=lambda x: x[1])

    features = []
    cat_features = []
    for line in reader:
        if file_type == 'tfrecord':
            line = parse_tfrecord(line)
        if example_ids is not None:
            example_ids.append(str(line['example_id']))
        if raw_ids is not None:
            raw_ids.append(str(line['raw_id']))
        if labels is not None:
            labels.append(float(line['label']))
        features.append([float(line[i]) for i in cont_columns])
        cat_features.append([int(line[i]) for i in cat_columns])

    features = np.array(features, dtype=np.float)
    cat_features = np.array(cat_features, dtype=np.int32)
    if labels is not None:
        labels = np.asarray(labels, dtype=np.float)

    return features, cat_features, cont_columns, cat_columns, \
        labels, example_ids, raw_ids


def read_data_dir(file_ext, file_type, path, require_example_ids,
                  require_labels, ignore_fields, cat_fields):
    if not tf.io.gfile.isdir(path):
        return read_data(
            file_type, path, require_example_ids,
            require_labels, ignore_fields, cat_fields)

    files = []
    for dirname, _, filenames in tf.io.gfile.walk(path):
        for filename in filenames:
            _, ext = os.path.splitext(filename)
            if file_ext and ext != file_ext:
                continue
            subdirname = os.path.join(path, os.path.relpath(dirname, path))
            files.append(os.path.join(subdirname, filename))

    features = None
    for fullname in files:
        ifeatures, icat_features, icont_columns, icat_columns, \
            ilabels, iexample_ids, iraw_ids = read_data(
                file_type, fullname, require_example_ids,
                require_labels, ignore_fields, cat_fields
            )
        if features is None:
            features = ifeatures
            cat_features = icat_features
            cont_columns = icont_columns
            cat_columns = icat_columns
            labels = ilabels
            example_ids = iexample_ids
            raw_ids = iraw_ids
        else:
            assert cont_columns == icont_columns, \
                "columns mismatch between files %s vs %s"%(
                    cont_columns, icont_columns)
            assert cat_columns == icat_columns, \
                "columns mismatch between files %s vs %s"%(
                    cat_columns, icat_columns)
            features = np.concatenate((features, ifeatures), axis=0)
            cat_features = np.concatenate(
                (cat_features, icat_features), axis=0)
            if labels is not None:
                labels = np.concatenate((labels, ilabels), axis=0)
            if example_ids is not None:
                example_ids.extend(iexample_ids)
            if raw_ids is not None:
                raw_ids.extend(iraw_ids)

    assert features is not None, "No data found in %s"%path

    return features, cat_features, cont_columns, cat_columns, \
        labels, example_ids, raw_ids


def train(args, booster):
    X, cat_X, X_names, cat_X_names, y, example_ids, _ = read_data_dir(
        args.file_ext, args.file_type, args.data_path, args.verify_example_ids,
        args.role != 'follower', args.ignore_fields, args.cat_fields)

    if args.validation_data_path:
        val_X, val_cat_X, val_X_names, val_cat_X_names, val_y, \
            val_example_ids, _ = \
            read_data_dir(
                args.file_ext, args.file_type, args.validation_data_path,
                args.verify_example_ids, args.role != 'follower',
                args.ignore_fields, args.cat_fields)
        assert X_names == val_X_names, \
            "Train data and validation data must have same features"
        assert cat_X_names == val_cat_X_names, \
            "Train data and validation data must have same features"
    else:
        val_X = val_cat_X = X_names = val_y = val_example_ids = None

    if args.output_path:
        tf.io.gfile.makedirs(os.path.dirname(args.output_path))
    if args.checkpoint_path:
        tf.io.gfile.makedirs(args.checkpoint_path)

    booster.fit(
        X, y,
        cat_features=cat_X,
        checkpoint_path=args.checkpoint_path,
        example_ids=example_ids,
        validation_features=val_X,
        validation_cat_features=val_cat_X,
        validation_labels=val_y,
        validation_example_ids=val_example_ids,
        output_path=args.output_path,
        feature_names=X_names,
        cat_feature_names=cat_X_names)


def write_predictions(filename, pred, example_ids=None, raw_ids=None):
    logging.debug("Writing predictions to %s.tmp", filename)
    headers = []
    lines = []
    if example_ids is not None:
        headers.append('example_id')
        lines.append(example_ids)
    if raw_ids is not None:
        headers.append('raw_id')
        lines.append(raw_ids)
    headers.append('prediction')
    lines.append(pred)
    lines = zip(*lines)

    fout = tf.io.gfile.GFile(filename+'.tmp', 'w')
    fout.write(','.join(headers) + '\n')
    for line in lines:
        fout.write(','.join([str(i) for i in line]) + '\n')
    fout.close()

    logging.debug("Renaming %s.tmp to %s", filename, filename)
    tf.io.gfile.rename(filename+'.tmp', filename, overwrite=True)

def test_one_file(args, bridge, booster, data_file, output_file):
    if data_file is None:
        X = cat_X = X_names = cat_X_names = y = example_ids = raw_ids = None
    else:
        X, cat_X, X_names, cat_X_names, y, example_ids, raw_ids = \
            read_data(
                args.file_type, data_file, args.verify_example_ids,
                False, args.ignore_fields, args.cat_fields)

    pred = booster.batch_predict(
        X,
        example_ids=example_ids,
        cat_features=cat_X,
        feature_names=X_names,
        cat_feature_names=cat_X_names)

    if y is not None:
        metrics = booster.loss.metrics(pred, y)
    else:
        metrics = {}
    logging.info("Test metrics: %s", metrics)

    if args.role == 'follower':
        bridge.start(bridge.new_iter_id())
        bridge.receive(bridge.current_iter_id, 'barrier')
        bridge.commit()

    if output_file:
        tf.io.gfile.makedirs(os.path.dirname(output_file))
        write_predictions(output_file, pred, example_ids, raw_ids)

    if args.role == 'leader':
        bridge.start(bridge.new_iter_id())
        bridge.send(
            bridge.current_iter_id, 'barrier', np.asarray([1]))
        bridge.commit()


class DataBlockLoader(object):
    def __init__(self, role, bridge, data_path, ext,
                 worker_rank=0, num_workers=1, output_path=None):
        self._role = role
        self._bridge = bridge
        self._num_workers = num_workers
        self._worker_rank = worker_rank
        self._output_path = output_path

        self._tm_role = 'follower' if role == 'leader' else 'leader'

        if data_path:
            files = None
            if not tf.io.gfile.isdir(data_path):
                files = [os.path.basename(data_path)]
                data_path = os.path.dirname(data_path)
            self._trainer_master = LocalTrainerMasterClient(
                self._tm_role, data_path, files=files, ext=ext)
        else:
            self._trainer_master = None

        self._count = 0
        if self._role == 'leader':
            self._block_queue = queue.Queue()
            self._bridge.register_data_block_handler(self._data_block_handler)
            self._bridge.start(self._bridge.new_iter_id())
            self._bridge.send(
                self._bridge.current_iter_id, 'barrier', np.asarray([1]))
            self._bridge.commit()
        elif self._role == 'follower':
            self._bridge.start(self._bridge.new_iter_id())
            self._bridge.receive(self._bridge.current_iter_id, 'barrier')
            self._bridge.commit()

    def _data_block_handler(self, msg):
        logging.debug('DataBlock: recv "%s" at %d', msg.block_id, msg.count)
        assert self._count == msg.count
        if not msg.block_id:
            block = None
        elif self._trainer_master is not None:
            block = self._trainer_master.request_data_block(msg.block_id)
            return False
        else:
            block = DataBlockInfo(msg.block_id, None)
        self._count += 1
        self._block_queue.put(block)
        return True

    def _request_data_block(self):
        while True:
            for _ in range(self._worker_rank):
                self._trainer_master.request_data_block()
            block = self._trainer_master.request_data_block()
            for _ in range(self._num_workers - self._worker_rank - 1):
                self._trainer_master.request_data_block()
            if block is None or self._output_path is None or \
                    not tf.io.gfile.exists(os.path.join(
                        self._output_path, block.block_id) + '.output'):
                break
        return block

    def get_next_block(self):
        if self._role == 'local':
            return self._request_data_block()

        if self._tm_role == 'leader':
            while True:
                block = self._request_data_block()
                if block is not None:
                    if not self._bridge.load_data_block(
                            self._count, block.block_id):
                        continue
                else:
                    self._bridge.load_data_block(self._count, '')
                break
            self._count += 1
        else:
            block = self._block_queue.get()
        return block


def test(args, bridge, booster):
    if not args.no_data:
        assert args.data_path, "Data path must not be empty"
    else:
        assert not args.data_path and args.role == 'leader'

    data_loader = DataBlockLoader(
        args.role, bridge, args.data_path, args.file_ext,
        args.worker_rank, args.num_workers, args.output_path)

    while True:
        data_block = data_loader.get_next_block()
        if data_block is None:
            break
        if args.output_path:
            output_file = os.path.join(
                args.output_path, data_block.block_id) + '.output'
        else:
            output_file = None
        test_one_file(
            args, bridge, booster, data_block.data_path, output_file)


def run(args):
    if args.verbosity == 0:
        logging.basicConfig(level=logging.WARNING)
    elif args.verbosity == 1:
        logging.basicConfig(level=logging.INFO)
    else:
        logging.basicConfig(level=logging.DEBUG)
    
    assert args.role in ['leader', 'follower', 'local'], \
        "role must be leader, follower, or local"
    assert args.mode in ['train', 'test', 'eval'], \
        "mode must be train, test, or eval"
    #follower或leader
    if args.role != 'local':
        bridge = Bridge(args.role, int(args.local_addr.split(':')[1]),
                        args.peer_addr, args.application_id, 0,
                        streaming_mode=args.use_streaming)
    else:
        bridge = None

    try:
        #boost
        booster = BoostingTreeEnsamble(
            bridge,
            learning_rate=args.learning_rate,
            max_iters=args.max_iters,
            max_depth=args.max_depth,
            l2_regularization=args.l2_regularization,
            max_bins=args.max_bins,
            num_parallel=args.num_parallel,
            loss_type=args.loss_type,
            send_scores_to_follower=args.send_scores_to_follower,
            send_metrics_to_follower=args.send_metrics_to_follower)
        #加载已存储的模型
        if args.load_model_path:
            booster.load_saved_model(args.load_model_path)
        #训练不需要bridge，为什么呢
        if args.mode == 'train':
            train(args, booster)
        #测试，评估模型需要bridge
        else:  # args.mode == 'test, eval'
            test(args, bridge, booster)
        #把模型存起来
        if args.export_path:
            booster.save_model(args.export_path)
    except Exception as e:
        logging.fatal(
            'Exception raised during training: %s',
            traceback.format_exc())
        raise e
    finally:
        #结束bridge
        if bridge:
            bridge.terminate()


if __name__ == '__main__':
    run(create_argument_parser().parse_args())
