# -*- coding: utf-8 -*-
import os
import psutil
import traceback

import h5py
import numpy as np

from buffalo.data import prepro
from buffalo.misc import aux, log
from buffalo.data.base import Data, DataOption


class MatrixMarketOptions(DataOption):
    def get_default_option(self) -> aux.Option:
        opt = {
            'type': 'matrix_market',
            'input': {
                'main': '',
                'uid': '',  # if not set, row-id is used as userid.
                'iid': ''  # if not set, col-id is used as itemid.
            },
            'data': {
                'internal_data_type': 'matrix',
                'validation': {
                    'name': 'sample',
                    'p': 0.01,
                    'max_samples': 500
                },
                'batch_mb': 1024,
                'use_cache': False,
                'tmp_dir': '/tmp/',
                'path': './mm.h5py'
            }
        }
        return aux.Option(opt)

    def is_valid_option(self, opt) -> bool:
        assert super(MatrixMarketOptions, self).is_valid_option(opt)
        if not opt['type'] == 'matrix_market':
            raise RuntimeError('Invalid data type: %s' % opt['type'])
        if opt['data']['internal_data_type'] != 'matrix':
            raise RuntimeError('MatrixMarket only support internal data type(matrix)')
        return True


class MatrixMarket(Data):
    def __init__(self, opt, *args, **kwargs):
        super().__init__(opt, *args, **kwargs)
        self.name = 'MatrixMarket'
        self.logger = log.get_logger('MatrixMarket')
        if isinstance(self.value_prepro,
                      (prepro.SPPMI)):
            raise RuntimeError(f'{self.opt.data.value_prepro.name} does not support MatrixMarket')
        self.data_type = 'matrix'

    def _create(self, data_path, P, H):
        def get_max_column_length(fname):
            with open(fname) as fin:
                max_col = 0
                for l in fin:
                    max_col = max(max_col, len(l))
            return max_col

        uid_path, iid_path, main_path = P['uid_path'], P['iid_path'], P['main_path']
        num_users, num_items, num_nnz = map(int, H.split())
        # Manually updating progress bar is a bit naive
        with log.pbar(log.DEBUG, total=5, mininterval=30) as pbar:
            uid_max_col = len(str(num_users)) + 1
            if uid_path:
                uid_max_col = get_max_column_length(uid_path) + 1
            pbar.update(1)
            iid_max_col = len(str(num_items)) + 1
            if iid_path:
                iid_max_col = get_max_column_length(iid_path) + 1
            pbar.update(1)
            try:
                db = self._create_database(data_path,
                                           num_users=num_users,
                                           num_items=num_items,
                                           num_nnz=num_nnz,
                                           uid_max_col=uid_max_col,
                                           iid_max_col=iid_max_col)
                idmap = db['idmap']
                # if not given, assume id as is
                if uid_path:
                    with open(uid_path) as fin:
                        idmap['rows'][:] = np.loadtxt(fin, dtype=f'S{uid_max_col}')
                else:
                    idmap['rows'][:] = np.array([str(i) for i in range(1, num_users + 1)],
                                                dtype=f'S{uid_max_col}')
                pbar.update(1)
                if iid_path:
                    with open(iid_path) as fin:
                        idmap['cols'][:] = np.loadtxt(fin, dtype=f'S{iid_max_col}')
                else:
                    idmap['cols'][:] = np.array([str(i) for i in range(1, num_items + 1)],
                                                dtype=f'S{iid_max_col}')
                pbar.update(1)
                num_header_lines = 0
                with open(main_path) as fin:
                    for line in fin:
                        if line.strip().startswith('%'):
                            num_header_lines += 1
                        else:
                            break
                pbar.update(1)
            except Exception as e:
                self.logger.error('Cannot create db: %s' % (str(e)))
                self.logger.error(traceback.format_exc())
                raise
        return db, num_header_lines

    def _create_working_data(self, db, source_path, ignore_lines):
        """
        Args:
            source_path: source data file path
            ignore_lines: number of lines to skip from start line
        """
        vali_indexes = [] if 'vali' not in db else db['vali']['indexes']
        vali_lines = []
        file_path = aux.get_temporary_file(self.opt.data.tmp_dir)
        with open(file_path, 'w') as w:
            fin = open(source_path, mode='r')
            file_size = fin.seek(0, 2)
            fin.seek(0, 0)
            for _ in range(ignore_lines):
                fin.readline()
            total = file_size - fin.tell()
            buffered = ''
            CHUNK_SIZE = 4096 * 1000
            total_lines = 0
            vali_indexes = sorted(vali_indexes)
            target_index = vali_indexes[0] if vali_indexes else -1
            vali_indexes = vali_indexes[1:]
            with log.pbar(log.INFO, total=total, mininterval=10) as pbar:
                while True:
                    buffered += fin.read(CHUNK_SIZE)
                    if buffered == '':
                        break
                    current_file_position = fin.tell()
                    pbar.update(CHUNK_SIZE)
                    num_lines_on_buffer = buffered.count('\n')
                    # search the position of validation sample and extract
                    # it from training data
                    while target_index >= 0 and target_index <= (total_lines + num_lines_on_buffer):
                        no_line = total_lines
                        new_buffered = ''
                        from_index = 0
                        for idx, c in enumerate(buffered):
                            if c == '\n':
                                if no_line == target_index:
                                    vali_lines.append(buffered[from_index:idx])
                                    if from_index > 0:
                                        w.write(buffered[0:from_index])
                                    new_buffered = buffered[idx + 1:]
                                    no_line += 1
                                    total_lines += 1
                                    num_lines_on_buffer -= 1
                                    break
                                no_line += 1
                                total_lines += 1
                                from_index = idx + 1
                                num_lines_on_buffer -= 1
                        buffered = new_buffered
                        if vali_indexes:
                            target_index, vali_indexes = vali_indexes[0], vali_indexes[1:]
                        else:
                            target_index = -1
                    where = buffered.rfind('\n')
                    total_lines += num_lines_on_buffer
                    if where != -1:
                        w.write(buffered[:where + 1])
                        buffered = buffered[where + 1:]
                    elif current_file_position == file_size:
                        w.write(buffered)
                        buffered = ''
            w.close()
            fin.close()
            return w.name, vali_lines

    def create(self) -> h5py.File:
        mm_main_path = self.opt.input.main
        mm_uid_path = self.opt.input.uid
        mm_iid_path = self.opt.input.iid

        data_path = self.opt.data.path
        if os.path.isfile(data_path) and self.opt.data.use_cache:
            self.logger.info('Use cached DB on %s' % data_path)
            self.open(data_path)
            return

        self.logger.info('Create the database from matrix market file.')
        with open(mm_main_path) as fin:
            header = '%'
            while header.startswith('%'):
                header = fin.readline()
            self.logger.debug('Building meta part...')
            db, num_header_lines = self._create(data_path,
                                                {'main_path': mm_main_path,
                                                 'uid_path': mm_uid_path,
                                                 'iid_path': mm_iid_path},
                                                header)
            try:
                num_header_lines += 1  # add metaline
                self.logger.info('Creating working data...')
                tmp_main, validation_data = self._create_working_data(db,
                                                                      mm_main_path,
                                                                      num_header_lines)
                self.logger.debug(f'Working data is created on {tmp_main}')
                self.logger.info('Building data part...')
                self._build_data(db, tmp_main, validation_data)
                db.attrs['completed'] = 1
                db.close()
                self.handle = h5py.File(data_path, 'r')
                self.path = data_path
            except Exception as e:
                self.logger.error('Cannot create db: %s' % (str(e)))
                self.logger.error(traceback.format_exc().splitlines())
                raise
            finally:
                if hasattr(self, 'patr'):
                    if os.path.isfile(self.path):
                        os.remove(self.path)
        self.logger.info('DB built on %s' % data_path)
