import numpy as np
from sklearn.tree._tree import Tree
from sklearn.tree import DecisionTreeRegressor

import warnings
warnings.filterwarnings("ignore")


def _version_range_check(version, lower, upper):
    return lower <= version <= upper


def _get_extended_nodes(nodes, version_in):
    if '0.21.3' <= version_in < '1.3.0':
        return [node + (0,) for node in nodes]
    return nodes


def _build_dtype_dict(dtypes, version_in):
    field_names = dtypes.names
    formats = [dtypes.fields[name][0] for name in field_names]
    offsets = [dtypes.fields[name][1] for name in field_names]
    itemsize = dtypes.itemsize

    if '0.21.3' <= version_in < '1.3.0':
        return {
            'field_names': list(field_names + ('missing_go_to_left',)),
            'formats': [str(fmt) for fmt in formats + [np.dtype('uint8')]],
            'offsets': [int(off) for off in offsets + [56]],
            'itemsize': 64
        }

    return {
        'field_names': list(field_names),
        'formats': [str(fmt) for fmt in formats],
        'offsets': [int(off) for off in offsets],
        'itemsize': int(itemsize)
    }


def _get_metadata(model, version_in):
    if _version_range_check(version_in, '0.21.3', '0.22.1'):
        return {
            'n_features_in': None,
            'n_features': model.n_features_,
            'n_outputs': model.n_outputs_,
            'n_classes': model.n_classes_
        }
    elif _version_range_check(version_in, '0.23.0', '0.23.2'):
        return {
            'n_features_in': model.n_features_in_,
            'n_features': model.n_features_,
            'n_outputs': model.n_outputs_,
            'n_classes': model.n_classes_
        }
    elif _version_range_check(version_in, '0.24.0', '1.1.3'):
        return {
            'n_features_in': model.n_features_in_,
            'n_features': model.n_features_,
            'n_outputs': model.n_outputs_,
            'n_classes': None
        }
    elif version_in >= '1.2.0':
        return {
            'n_features_in': model.n_features_in_,
            'n_features': None,
            'n_outputs': model.n_outputs_,
            'n_classes': None
        }
    return {}


def serialize_decision_tree_reg(model, version_in):
    tree = model.tree_
    state = tree.__getstate__()

    serialized_tree = {
        'max_depth': int(state['max_depth']),
        'node_count': int(state['node_count']),
        'values': state['values'].tolist(),
        'nodes': [list(n) for n in _get_extended_nodes(state['nodes'].tolist(), version_in)],
        'dtypes': _build_dtype_dict(state['nodes'].dtype, version_in)
    }

    metadata = _get_metadata(model, version_in)
    metadata['serialized_tree'] = serialized_tree
    metadata['version_sklearn_in'] = version_in

    return metadata


def _build_tree_dtype(dtypes_dict, version_out):
    version_lt_1_3 = version_out < '1.3.0'
    num_elements = 7 if version_lt_1_3 else 8

    field_names = dtypes_dict['field_names'][:num_elements]
    formats = [np.dtype(fmt) for fmt in dtypes_dict['formats'][:num_elements]]
    offsets = dtypes_dict['offsets'][:num_elements]
    itemsize = 56 if version_lt_1_3 else 64

    return np.dtype({
        'names': field_names,
        'formats': formats,
        'offsets': offsets,
        'itemsize': itemsize
    }), num_elements


def deserialize_decision_tree_reg(data, version_out):
    version_in = data['version_sklearn_in']
    serialized = data['serialized_tree']
    dtype_dict = serialized['dtypes']

    tree_dtype, num_elements = _build_tree_dtype(dtype_dict, version_out)

    serialized['nodes'] = [tuple(n[:num_elements]) for n in serialized['nodes']]
    nodes_array = np.array(serialized['nodes'], dtype=tree_dtype)
    values_array = np.array(serialized['values'])

    n_classes = np.array([1], dtype=np.intp)  # regression
    n_outputs = data['n_outputs']
    n_features = (data['n_features'] or data['n_features_in'])

    tree_obj = Tree(n_features, n_classes, n_outputs)
    tree_obj.__setstate__({
        'max_depth': serialized['max_depth'],
        'node_count': serialized['node_count'],
        'nodes': nodes_array,
        'values': values_array
    })

    new_tree = DecisionTreeRegressor(max_depth=serialized['max_depth'], random_state=42)
    new_tree.tree_ = tree_obj
    new_tree.n_outputs_ = n_outputs

    if _version_range_check(version_out, '0.21.3', '0.22.1'):
        new_tree.n_features_ = n_features
    elif _version_range_check(version_out, '0.23.0', '0.24.2'):
        new_tree.n_features_ = n_features
        new_tree.n_features_in_ = n_features
    elif version_out >= '1.0.0':
        new_tree.n_features_in_ = n_features

    return new_tree