# -*- coding: utf-8 -*-
"""
Created on 2017-12-4

@author: cheng.li
"""

from typing import List
import numpy as np
from distutils.version import LooseVersion
from sklearn import __version__ as sklearn_version
from sklearn.ensemble import RandomForestRegressor as RandomForestRegressorImpl
from xgboost import XGBRegressor as XGBRegressorImpl
from alphamind.model.modelbase import ModelBase
from alphamind.utilities import alpha_logger


class RandomForestRegressor(ModelBase):

    def __init__(self, n_estimators: int=100, features: List=None, **kwargs):
        super().__init__(features)
        self.impl = RandomForestRegressorImpl(n_estimators, **kwargs)
        self.trained_time = None

    def score(self, x: np.ndarray, y: np.ndarray) -> float:
        return self.impl.score(x, y)

    def save(self) -> dict:
        model_desc = super().save()
        model_desc['sklearn_version'] = sklearn_version

    @classmethod
    def load(cls, model_desc: dict):
        obj_layout = super().load(model_desc)

        if LooseVersion(sklearn_version) < LooseVersion(model_desc['sklearn_version']):
            alpha_logger.warning('Current sklearn version {0} is lower than the model version {1}. '
                                 'Loaded model may work incorrectly.'.format(
                sklearn_version, model_desc['sklearn_version']))
        return obj_layout


class XGBRegressor(ModelBase):

    def __init__(self,
                 n_estimators: int=100,
                 learning_rate: float=0.1,
                 max_depth: int=3,
                 features: List=None, **kwargs):
        super().__init__(features)
        self.impl = XGBRegressorImpl(n_estimators=n_estimators,
                                     learning_rate=learning_rate,
                                     max_depth=max_depth,
                                     **kwargs)



@@ -279,24 +279,26 @@
 gressorImpl%0A
+# 
 from xgboost
@@ -1371,16 +1371,18 @@
 ayout%0A%0A%0A
+# 
 class XG
@@ -1396,33 +1396,36 @@
 sor(ModelBase):%0A
-%0A
+#%0A# 
     def __init__
@@ -1431,16 +1431,17 @@
 _(self,%0A
+#
         
@@ -1449,16 +1449,17 @@
         
+ 
 n_estima
@@ -1473,16 +1473,18 @@
 nt=100,%0A
+# 
         
@@ -1518,16 +1518,18 @@
 at=0.1,%0A
+# 
         
@@ -1555,16 +1555,17 @@
  int=3,%0A
+#
         
@@ -1573,16 +1573,17 @@
         
+ 
 features
@@ -1598,32 +1598,34 @@
 one, **kwargs):%0A
+# 
         super().
@@ -1635,32 +1635,34 @@
 nit__(features)%0A
+# 
         self.imp
@@ -1709,16 +1709,17 @@
 mators,%0A
+#
         
@@ -1739,32 +1739,33 @@
                 
+ 
 learning_rate=le
@@ -1777,16 +1777,18 @@
 g_rate,%0A
+# 
         
@@ -1837,16 +1837,18 @@
 _depth,%0A
+# 
         

comment out XGBRegressor