import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

# 데이터 불러오기
df = pd.read_excel('imsi3.xlsx', sheet_name='Sheet1')

# 인구증감수 열 제거
#imsi3.xlsx : 연도,출생수,사망수,인구수
#선형회귀로 나온 데이터에다 인구유입등의 자료가 있으면 플러스 해줌

# 학습용 데이터와 테스트용 데이터로 분리
X = df.drop('인구수', axis=1) # 독립 변수
y = df['인구수'] # 종속 변수

#가지고 있는 데이터중 학습용데이터(80%)와 테스트용데이터(0.2라서 20%)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 선형회귀 모델 학습
lr = LinearRegression()
lr.fit(X_train, y_train)

# 테스트 데이터로 예측하기
#출생수와 사망수는 마지막자료인 2022년도 자료를 넣음
X_pred = pd.DataFrame({
'연도': [2023, 2024, 2025],
    '출생수': [14500, 14500, 14500],
    '사망수': [19000, 19000, 19000],
})

y_pred = lr.predict(X_pred)

# 예측 결과 출력
for year, pop in zip(X_pred['연도'], y_pred):
print(f'{year}년 인구수 예측: {pop:.0f}')