import pandas as pd
import numpy as np
import streamlit as st
import tensorflow as tf
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split
from tensorflow.keras.layers import Input, Embedding, Flatten, Concatenate, Dense, BatchNormalization, Dropout
from tensorflow.keras.models import Model
import matplotlib.pyplot as plt

from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.linear_model import LogisticRegression



# app.py

# Função de Preprocessamento
def preprocess_data(base_final):
    # Convertendo indenização e tempo para numéricos
    base_final['indenizacao_num'] = pd.to_numeric(base_final['indenizacao'], errors='coerce')
    base_final['tempo_num'] = pd.to_numeric(base_final['tempo_meses'], errors='coerce')

    # Identificar variáveis categóricas e numéricas
    categorical_cols = ['classe', 'comarca', 'danos', 'resultado']
    numerical_cols = base_final.select_dtypes(include=['float64', 'int64']).columns

    # Imputação para variáveis categóricas
    from sklearn.impute import SimpleImputer
    imputer_cat = SimpleImputer(strategy='most_frequent')
    base_final[categorical_cols] = imputer_cat.fit_transform(base_final[categorical_cols])

    # Imputação para variáveis numéricas
    imputer_num = SimpleImputer(strategy='mean')
    base_final[numerical_cols] = imputer_num.fit_transform(base_final[numerical_cols])

    # Criar dummies para todas as variáveis categóricas
    from sklearn.preprocessing import OneHotEncoder
    encoder = OneHotEncoder(drop='first', sparse_output=False)
    dummies = pd.DataFrame(
        encoder.fit_transform(base_final[categorical_cols]),
        columns=encoder.get_feature_names_out(categorical_cols),
        index=base_final.index
    )

    # Combinar variáveis contínuas com dummies
    base_final_preprocessed = pd.concat([base_final, dummies], axis=1)
    base_final_preprocessed.drop(columns=categorical_cols, inplace=True)
    base_final_preprocessed = base_final_preprocessed.select_dtypes(include=[np.number])

    return base_final_preprocessed

# Função de Modelagem para Indenização
def model_indenizacao(base_final_preprocessed):
    X = base_final_preprocessed.drop(columns=['tempo_num', 'indenizacao_num'], errors='ignore')
    y_indenizacao = base_final_preprocessed['indenizacao_num']

    # Divisão entre treino e teste
    X_train, X_test, y_train_indenizacao, y_test_indenizacao = train_test_split(X, y_indenizacao, test_size=0.3, random_state=42)

    rf_model_indenizacao = RandomForestRegressor(n_estimators=100, random_state=42)
    rf_model_indenizacao.fit(X_train, y_train_indenizacao)
    y_pred_indenizacao = rf_model_indenizacao.predict(X_test)
    rmse_indenizacao = np.sqrt(mean_squared_error(y_test_indenizacao, y_pred_indenizacao))

    return rf_model_indenizacao, rmse_indenizacao

# Função de Modelagem para Tempo
def model_tempo(base_final_preprocessed):
    X = base_final_preprocessed.drop(columns=['tempo_num', 'indenizacao_num'], errors='ignore')
    y_tempo = base_final_preprocessed['tempo_num']

    # Divisão entre treino e teste
    X_train, X_test, y_train_tempo, y_test_tempo = train_test_split(X, y_tempo, test_size=0.3, random_state=42)

    rf_model_tempo = RandomForestRegressor(n_estimators=100, random_state=42)
    rf_model_tempo.fit(X_train, y_train_tempo)
    y_pred_tempo = rf_model_tempo.predict(X_test)
    rmse_tempo = np.sqrt(mean_squared_error(y_test_tempo, y_pred_tempo))

    return rf_model_tempo, rmse_tempo

# Função de Classificação de Risco
def model_risco(base_final, X_test, rf_model_indenizacao, rf_model_tempo):
    logistic_model = LogisticRegression(max_iter=500)
    logistic_model.fit(X_test, (base_final['resultado'] == 'Procedente').astype(int))
    y_pred_proba = logistic_model.predict_proba(X_test)[:, 1]
    y_pred_indenizacao = logistic_model.predict(X_test)
    provisionamento_esperado = (y_pred_indenizacao * y_pred_proba).sum()

    # Classificação de Risco
    X_test_df = base_final.loc[X_test.index].copy()
    X_test_df['prob_perda'] = y_pred_proba
    X_test_df['risco'] = np.select(
        [X_test_df['prob_perda'] < 0.3, 
         (X_test_df['prob_perda'] >= 0.3) & (X_test_df['prob_perda'] < 0.7), 
         X_test_df['prob_perda'] >= 0.7],
        ['Remoto', 'Possível', 'Provável'],
        default='Desconhecido'
    )

    return provisionamento_esperado, X_test_df

# Título do App
st.title("Sistema Jurídico - Modelagem de Indenização e Tempo de Processo")

# Upload de arquivo
uploaded_file = st.file_uploader("Faça o upload do arquivo Excel", type=["xlsx"])

if uploaded_file:
    # Carregar os dados
    base_final = pd.read_excel(uploaded_file)

    st.subheader("Passo 1: Preprocessamento")
    base_final_preprocessed = preprocess_data(base_final)
    st.write("Dados preprocessados:", base_final_preprocessed.head(100))

    st.subheader("Passo 2: Modelos Preditivos")
    rf_model_indenizacao, rmse_indenizacao = model_indenizacao(base_final_preprocessed)
    st.write(f"RMSE (Indenização): {rmse_indenizacao:.2f}")

    rf_model_tempo, rmse_tempo = model_tempo(base_final_preprocessed)
    st.write(f"RMSE (Tempo do Processo): {rmse_tempo:.2f}")

    st.subheader("Passo 3: Provisionamento e Classificação de Risco")
    provisionamento_esperado, X_test_df = model_risco(base_final, base_final_preprocessed, rf_model_indenizacao, rf_model_tempo)
    st.write(f"Provisionamento esperado: R$ {provisionamento_esperado:.2f}")
    st.write("Classificação de Risco:", X_test_df[['risco', 'indenizacao_num']].head(100))

    # st.subheader("Passo 4: Classificação de Tempo do Processo")
    # X_test_df['tempo_estimado'] = rf_model_tempo.predict(X_test_df)
    # X_test_df['categoria_tempo'] = np.select(
    #     [X_test_df['tempo_estimado'] < 6, 
    #      (X_test_df['tempo_estimado'] >= 6) & (X_test_df['tempo_estimado'] < 12), 
    #      X_test_df['tempo_estimado'] >= 12],
    #     ['Curto', 'Moderado', 'Longo'],
    #     default='Desconhecido'
    # )
    # st.write("Classificação de Tempo:", X_test_df[['categoria_tempo', 'tempo_meses']].head(100))

    # Gráficos opcionais
    st.subheader("Distribuição de Previsões")
    st.bar_chart(X_test_df['risco'].value_counts())