import yfinance as yf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 1. 拉数据：贵州茅台最近 1 年
df = yf.download("600519.SS", start="2024-01-01", end="2025-01-01")
df.columns = [c.lower() for c in df.columns]

# 2. 用 5 日高低点找局部极值（zigzag 思路）
window = 5
df['hh'] = df['high'] == df['high'].rolling(window*2+1, center=True).max()
df['ll'] = df['low']  == df['low'].rolling(window*2+1, center=True).min()

# 3. 提取所有高点、低点
highs = df[df['hh']]['high'].tolist()
lows  = df[df['ll']]['low'].tolist()

# 4. 判断 HH / HL：滚动比较
hh_flag, last_h = [], -np.inf
for h in highs:
    hh_flag.append(h > last_h)
    if h > last_h:
        last_h = h

hl_flag, last_l = [], np.inf
for l in lows:
    hl_flag.append(l > last_l)
    if l < last_l:
        last_l = l

print(f"局部高点 {len(highs)} 个，其中 HH 数量: {sum(hh_flag)}")
print(f"局部低点 {len(lows)} 个，其中 HL 数量: {sum(hl_flag)}")

uptrend = sum(hh_flag) >= 3 and sum(hl_flag) >= 3
print("判定结果:", "上升趋势 ✅" if uptrend else "非上升趋势 ❌")

# 5. 画图
df['close'].plot(title="600519 贵州茅台", figsize=(12, 5))
plt.scatter(df[df['hh']].index, df[df['hh']]['high'], c='r', marker='^')
plt.scatter(df[df['ll']].index, df[df['ll']]['low'],  c='g', marker='v')
plt.show()

import yfinance as yf
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import torch, torch.nn as nn

# === 通用：拉数据 + 打标签 ===
df = yf.download("000300.SS", start="2018-01-01", end="2025-01-01")
df.columns = [c.lower() for c in df.columns]

# 标签：未来 20 日收益 > 5% 视为上升趋势
df['ret_20'] = df['close'].pct_change(20).shift(-20)
df['label'] = (df['ret_20'] > 0.05).astype(int)
df = df.dropna()

def features(d):
    out = pd.DataFrame()
    out['sma5']  = d['close'] / d['close'].rolling(5).mean()  - 1
    out['sma20'] = d['close'] / d['close'].rolling(20).mean() - 1
    out['slope'] = d['close'].rolling(10).apply(
        lambda x: np.polyfit(range(10), x, 1)[0])
    out['vol']   = d['close'].rolling(10).std()
    out['hh']    = (d['high'] > d['high'].shift(5)).astype(int).rolling(10).mean()
    out['hl']    = (d['low']  > d['low'].shift(5)).astype(int).rolling(10).mean()
    return out

X = features(df)
data = pd.concat([X, df['label']], axis=1).dropna()

# === 路线 A：随机森林 ===
X_train, X_test, y_train, y_test = train_test_split(
    data.drop('label', axis=1), data['label'], test_size=0.3, shuffle=False)
rf = RandomForestClassifier(n_estimators=200, max_depth=8, random_state=42)
rf.fit(X_train, y_train)
print("随机森林:")
print(classification_report(y_test, rf.predict(X_test)))

# === 路线 B：LSTM 序列模型 ===
SEQ = 20
arr = X.values
def make_seq(a, y, s=SEQ):
    Xs, ys = [], []
    for i in range(len(a)-s):
        Xs.append(a[i:i+s]); ys.append(y[i+s])
    return np.array(Xs), np.array(ys)

Xs, ys = make_seq(arr, data['label'].values)
split = int(len(Xs)*0.7)
Xtr = torch.tensor(Xs[:split], dtype=torch.float32)
ytr = torch.tensor(ys[:split], dtype=torch.long)
Xte = torch.tensor(Xs[split:], dtype=torch.float32)
yte = torch.tensor(ys[split:], dtype=torch.long)

class TrendLSTM(nn.Module):
    def __init__(self, n_feat, hidden=32):
        super().__init__()
        self.lstm = nn.LSTM(n_feat, hidden, batch_first=True)
        self.fc   = nn.Linear(hidden, 2)
    def forward(self, x):
        _, (h, _) = self.lstm(x)
        return self.fc(h[-1])

model = TrendLSTM(Xtr.shape[2])
opt   = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(15):
    model.train()
    perm = torch.randperm(len(Xtr))
    for i in range(0, len(Xtr), 64):
        idx = perm[i:i+64]
        out  = model(Xtr[idx])
        loss = loss_fn(out, ytr[idx])
        opt.zero_grad(); loss.backward(); opt.step()

model.eval()
pred = model(Xte).argmax(dim=1).numpy()
print("LSTM:")
print(classification_report(yte.numpy(), pred))