#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: sunluelectric@github
This program demonstrates using sklearn to solve a linear regression problem.
It also calculates the distance of the fitting curve to the origin.
"""
import numpy as np
from sklearn.linear_model import LinearRegression
import self_error

class LinearRegressionDemo:
    """
    LinearRegressionDemo solves a linear regression problem using sklearn.
    """
    def __init__(self, aaf_sample):
        self.aaf_sample = aaf_sample
        self.i_number_of_sample = len(self.aaf_sample)
        af_x = np.array([])
        af_y = np.array([])
        for iter_sample in self.aaf_sample:
            af_x = np.append(af_x, iter_sample[0])
            af_y = np.append(af_y, iter_sample[1])
        self.af_x = af_x.reshape((-1,1))
        self.af_y = af_y
        self.af_xinv = af_x
        self.af_yinv = af_y.reshape((-1,1))
        self.f_k = None
        self.f_b = None
        self.f_penalty = None
        self.f_kinv = None
        self.f_binv = None
        self.f_penalty_inv = None

    def solvelinearregression(self):
        """
        Use sklearn to solve the linear regression problem.
        """
        lr_model = LinearRegression().fit(self.af_x, self.af_y)
        self.f_k = lr_model.coef_[0]
        self.f_b = lr_model.intercept_
        self.f_penalty = np.sum(
            (self.af_y - (self.f_k * self.af_x + self.f_b))**2)
    def solvelinearregressioninv(self):
        """
        Use sklearn to solve the linear regression problem.
        """
        lr_model = LinearRegression().fit(self.af_yinv, self.af_xinv)
        self.f_kinv = lr_model.coef_[0]
        self.f_binv = lr_model.intercept_
        self.f_penalty_inv = np.sum(
            (self.af_xinv - (self.f_kinv * self.af_yinv + self.f_binv))**2)
    def calculatedistance(self):
        """
        calculatedistance calculates the distance (and direction)
        of the LR result to the origin.
        """
        if self.i_number_of_sample == 0:
            self_error.GeneralErrorMessage(
                "There is no input samples to the LR problem.")
            return np.array([0,0])
        if self._checkallequal():
            return self.aaf_sample[0]
        self.solvelinearregression()
        self.solvelinearregressioninv()
        if self.f_penalty <= self.f_penalty_inv:
            return np.array(
                [-(self.f_b*self.f_k)/(self.f_k**2+1),
                 self.f_b/(self.f_k**2+1)])
        return np.array(
            [self.f_binv/(self.f_kinv**2+1),
             -(self.f_binv*self.f_kinv)/(self.f_kinv**2+1)])
    def _checkallequal(self):
        """
        Check whether all samples are identical.
        """
        if self.i_number_of_sample == 1:
            return True
        b_flag = True
        for iter_sample in self.aaf_sample:
            if not np.array_equal(self.aaf_sample[0], iter_sample):
                b_flag = False
                break
        return b_flag
    def _calculateaverage(self):
        average_vector = np.zeros((1,2))
        for iter_sample in self.aaf_sample:
            average_vector += iter_sample
        average_vector = average_vector/self.i_number_of_sample
