#Datawhale 暑期AI夏令营task1 baseline代码

import numpy as np
import pandas as pd
from pathlib import Path
from sklearn.linear_model import LinearRegression

# 定义数据存储的基础路径，确保所有数据都存放在同级的data目录下
base_path = Path("data")

# 读取电力市场价格数据
electricity_price = pd.read_csv(base_path / "electricity price.csv")
# 读取发电机组数据
unit = pd.read_csv(base_path / "unit.csv")

# 准备示例提交数据 sample_submit
# 找到出清价格为缺失值的行，即需要预测的目标
sample_submit = electricity_price[electricity_price["clearing price (CNY/MWh)"].isna()].drop(columns="demand")
# 将示例提交数据保存到 CSV 文件中
sample_submit.to_csv(base_path / "sample_submit.csv", index=False)

# 将 day 和 time 列合并成 timestamp 列，便于提取时间戳特征
electricity_price["timestamp"] = pd.to_datetime(
    electricity_price["day"] + " " + electricity_price["time"].str.replace("24:00:00", "00:00"))

# 处理时间为 24:00:00 的情况，即表示第二天的 00:00:00
mask = electricity_price['timestamp'].dt.time == pd.Timestamp('00:00:00').time()
# 将这些行的日期部分加一天
electricity_price.loc[mask, 'timestamp'] += pd.Timedelta(days=1)

# 设置列的顺序，同时去除 day 和 time 列
electricity_price = electricity_price[["timestamp", "demand", "clearing price (CNY/MWh)"]]

# 显示前5行电力市场价格数据
print(electricity_price.head())

# 显示前5行发电机组数据
print(unit.head())

# 按照一度电的耗煤量（近似为边际成本）升序排序
sorted_unit = unit.sort_values("coal consumption (g coal/KWh)")
# 显示排序后的前5行发电机组数据
print(sorted_unit.head())

# 计算 sorted_unit 的累积容量
sorted_unit['cumulative_capacity'] = sorted_unit['Capacity（MW）'].cumsum()

# 初始化用于存储预测价格的列表
prices = []

# 根据需求找到最后一个满足总需求的机组报价
for demand in electricity_price["demand"]:
    price = sorted_unit[sorted_unit['cumulative_capacity'] >= demand]["coal consumption (g coal/KWh)"].iloc[0]
    prices.append(price)

# 输出预测价格的数量
print(len(prices))
# 显示前5个预测价格
print(prices[:5])

# 创建线性回归模型
model = LinearRegression()

# 定义训练集的长度
train_length = 55392

# 将预测价格转换为 NumPy 数组并调整其形状
prices = np.array(prices).reshape(-1, 1)
X = prices[:train_length]
# 获取训练集对应的实际出清价格
y = electricity_price["clearing price (CNY/MWh)"].iloc[:train_length].values.reshape(-1, 1)

# 训练线性回归模型
model.fit(X, y)

# 输出模型的系数和截距
print(model.coef_, model.intercept_)

# 使用模型对剩余数据进行预测
y_pred = model.predict(prices[train_length:])
# 将预测结果从2维矩阵转为1维数组
y_pred = y_pred.flatten()

# 显示前5个预测结果
print(y_pred[:5])

# 将预测的出清价格填入 sample_submit 数据框中
sample_submit["clearing price (CNY/MWh)"] = y_pred

# 显示前5行 sample_submit 数据
print(sample_submit.head())

# 将最终的提交数据保存到 CSV 文件中
sample_submit.to_csv("submit.csv", index=False)