# -*- coding: utf-8 -*-
"""
Created on Wed Mar 24 15:22:14 2021

@author: Adm
"""
import pandas as pd
import statsmodels.formula.api as smf
import matplotlib.pyplot as plt
%matplotlib qt # plotar na janela adicional

# Importar arquivo do Excel
# mudar diretório de acordo com o caminho no seu computador
# arq_tabela = r'G:\Meu Drive\Drive\Aulas\Algoritmos\dados_estr.xlsx'
# Se não instalado: pip install openpyxl
# dados = pd.read_excel(arq_tabela, engine = 'openpyxl')

# Se achar conveniente pode fazer o download dos dados
arq_xlxs = r'https://gmarcatti.github.io/dados/dados_estr.csv'

# Importar direto da web
dados = pd.read_table(arq_xlxs, decimal = ',', sep = ';')


# Regressão no python
# LinearRegression: sklearn
# ols: statsmodels
# curve_fit: scipy

# Ajustar
ml = smf.ols('vol ~ idade', data = dados).fit()
ml.summary()
ml.params
ml.params[0]
ml.params[1]
ml.params['Intercept']
ml.params['idade']
ml.rsquared

# Identificar os grupos na data frame
grupos = dados.groupby('estrato')

# Ajustar por grupos
coef = []
for g, grp in grupos:
    ml_grp = smf.ols('vol ~ idade', data = grp).fit()
    coef.append([g, ml_grp.params['Intercept'], ml_grp.params['idade']])

df_coef = pd.DataFrame(coef, columns = ['estrato', 'b0', 'b1'])

# Juntar os coeficientes dos ajustes com a base de dados
dados_coef = pd.merge(dados, df_coef, on = 'estrato')

# Gerar gráficos por grupos
for g, grp in grupos:
    plt.plot(grp['idade'], grp['vol'], 'o', label = g)

i = 0
for g, grp in grupos:
    vol_est = df_coef.at[i, 'b0'] + df_coef.at[i, 'b1'] * grp['idade']
    plt.plot(grp['idade'], vol_est, lw = 2, label = g)
    i += 1

plt.legend()