import os
from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance
from qdrant_client.models import PointStruct
from m98_rag import embd, readChunks, rerank

QDRANT_URL = os.getenv('QDRANT_URL', 'http://localhost:6333')
QDRANT_KEY = os.getenv('QDRANT_KEY', '')

client = QdrantClient(url=QDRANT_URL,
                      api_key=QDRANT_KEY,
                      timeout=100)


class Hits(list):
    def __init__(self, hits, query_: str, limit_: int):
        super().__init__(hits)
        self.query = query_
        self.limit = limit_
        self.res = None

    def rerank(self):
        if self.res is None:
            hits = [hit.payload['text'] for hit in self]
            self.res = rerank(hits, query_=self.query, top_n_=self.limit)
        return self.res

    def top(self, top_n=3, related_=True):
        res = self.rerank()
        if related_:
            return [hit[0] for hit in res[:top_n] if hit[1] > 0]
        else:
            return [hit[0] for hit in res[:top_n]]

    def print(self):
        return ('\n' + '=' * 50 + '\n').join(self.top())


class RAG:
    def __init__(self, collection_name="my_collection", vectors_size=768):
        self.collection_name = collection_name
        self.vectors_size = vectors_size
        if not client.collection_exists(self.collection_name):
            client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(size=self.vectors_size, distance=Distance.COSINE)
            )

    def upsert(self, chunks_: list, ID_: int = 0):
        for i in range(0, len(chunks_), 8):
            batch = chunks_[i: i + 8]
            print(batch)
            vectors = embd(batch)
            client.upsert(
                collection_name=self.collection_name,
                points=[
                    PointStruct(
                        id=ID_ + i + idx,
                        vector=vector[1],
                        payload={"text": vector[0]}
                    )
                    for idx, vector in enumerate(vectors)
                ]
            )

    def size(self):
        count = client.count(collection_name=self.collection_name)
        print(self.collection_name, 'size:', count.count)
        return count.count

    def clear(self):
        client.delete_collection(collection_name=self.collection_name)
        client.create_collection(
            collection_name=self.collection_name,
            vectors_config=VectorParams(size=self.vectors_size, distance=Distance.COSINE)
        )

    def search(self, query_: str, limit_: int = 5):
        query_vector = embd([query_])
        hits = client.search(
            collection_name=self.collection_name,
            query_vector=query_vector[0][1],
            limit=limit_  # Return 5 closest points
        )
        return Hits(hits, query_, limit_)


if __name__ == '__main__':
    print(QDRANT_URL)
    print(QDRANT_KEY)
    test = RAG()
    if test.size() <= 0:
        chunks = readChunks('./test.md')
        test.upsert(chunks)
    tmp = test.search('机器人限拥令是什么')
    print(tmp.print())

class StreamingLLM(Llama):
    pass
    def kv_cache_seq_trim(self):
        self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1)

    def kv_cache_seq_ltrim(self, n_keep, n_discard=256, n_past=-1):
        if n_past < 0:
            n_past = self.n_tokens
        self._ctx.kv_cache_seq_rm(-1, n_keep, n_keep + n_discard)
        self._ctx.kv_cache_seq_shift(0, n_keep + n_discard, n_past, -n_discard)
        self.input_ids[n_keep:n_past - n_discard] = self.input_ids[n_keep + n_discard:n_past]
        self.n_tokens = n_past - n_discard

    def _venv_init(self):
        self.venv = [0]
        self.venv_idx_map = []

    def venv_create(self, name: str):
        self.venv.append(0)
        self.venv_idx_map.append(name)
        return name

    def venv_disband(self, name_set):
        if len(self.venv) <= 1:
            return False
        name_set = {x for x in name_set if x in self.venv_idx_map}
        if not name_set:
            return False
        while self.venv_idx_map:
            if self.venv_idx_map[0] in name_set:
                self.venv_idx_map.pop(0)  # 删除
                tmp = self.venv.pop(1)  # 对应的 venv 移入上一层
                self.venv[0] += tmp
            else:
                break
        return True

    def venv_revision(self, name: str):
        if len(self.venv) <= 1:
            return False
        if name not in self.venv_idx_map:
            return False
        _s = 0
        while self.venv_idx_map:
            if self.venv_idx_map[-1] == name:
                break
            self.venv_idx_map.pop()  # 删除
            _s += self.venv.pop()
        if _s:
            self.n_tokens -= min(_s, self.n_tokens)
            self.kv_cache_seq_trim()
        return True

    def venv_remove(self, name: str):
        if len(self.venv) <= 1:
            return False
        if name not in self.venv_idx_map:
            return False
        venv_idx = self.venv_idx_map.index(name) + 1
        while self.venv_idx_map:
            self.venv_idx_map.pop(venv_idx - 1)  # 删除
            if venv_idx == len(self.venv) - 1:
                # 最后一层
                self.n_tokens -= min(self.venv.pop(), self.n_tokens)
                self.kv_cache_seq_trim()
                break
            else:
                # 非最后一层
                n_keep = self.n_tokens - sum(self.venv[i] for i in range(venv_idx, len(self.venv)))
                n_discard = self.venv.pop(venv_idx)
                self.kv_cache_seq_ltrim(n_keep, n_discard)
                try:
                    venv_idx = self.venv_idx_map.index(name, venv_idx - 1) + 1
                except ValueError:  # 没有了
                    break
        return True

    def eval_t(self, tokens, n_keep=4, n_discard=256, im_start=None):
        if self._n_ctx < self.n_tokens + len(tokens):
            tmp_n_discard = max(n_discard, self.n_tokens + len(tokens) - self._n_ctx)
            self.kv_cache_seq_ltrim(n_keep, tmp_n_discard)
        for i in range(0, len(tokens), self.n_batch):
            pass
            self.n_tokens += n_tokens
            self.venv[-1] += n_tokens

from transformers import AutoModelForCausalLM
import torch

# 提前将 modeling_bluelm.py 中用到 flash_attn 的部分改成 None，反正不真运行，只需要模型结构
tmp = AutoModelForCausalLM.from_pretrained(r'D:\models\BlueLM-7B',
                                     torch_dtype=torch.bfloat16,
                                     trust_remote_code=True)

test_i = torch.arange(0, 10, dtype=torch.long)

embedding = tmp.model.embed_tokens
layer_norm = tmp.model.embed_layer_norm

test_o_o = embedding(test_i)
test_o_o = layer_norm(test_o_o)

for param in embedding.parameters():
    if len(param.shape) > 1:
        param.data = layer_norm(param.data)

test_o_c = embedding(test_i)

print(torch.allclose(test_o_o, test_o_c, atol=1e-4))

del tmp.model.embed_layer_norm
tmp.save_pretrained(r'D:\models\BlueLM')
# 记得将缺失的一些文件手动复制一下
# 顺便删掉config.json里的rope scaling type

import os

class Fileset(list):
    def __init__(self, path, ext='', _read=None):
        if isinstance(path, str):
            self.root = path
            self.extend(f for f in os.listdir(self.root) if f.endswith(ext))
            self._read = _read

    def __getitem__(self, index):
        if isinstance(index, int):  # index是索引
            if self._read:
                return self._read(os.path.join(self.root, super().__getitem__(index)))
            else:
                return os.path.join(self.root, super().__getitem__(index))
        else:  # index是切片
            fileset = Fileset(None)
            fileset.root = self.root
            fileset._read = self._read
            fileset.extend(super().__getitem__(index))
            return fileset

    def getFileName(self, index):
        fname, ext = os.path.splitext(super().__getitem__(index))
        return fname


from tokenizer import tokenizer
token_eos = 2


def readOne(filePath):
    retn = []
    with open(file=filePath, encoding='utf-8') as f:
        for line in f:
            retn += tokenizer.encode(line).ids
    retn.append(token_eos)
    return retn


class Hcorpus():
    def __init__(self, path, ext='txt', fileset_idx=0, fileset_sub_idx=0):
        self.fileset = Fileset(path, ext, readOne)
        self.fileset_idx = fileset_idx
        self.fileset_sub_idx = fileset_sub_idx
        if self.fileset_sub_idx < 0:  # 再读上一个太复杂了，直接放弃
            self.fileset_sub_idx = 0
        if self.fileset_idx >= len(self.fileset):
            self.fileset_idx = 0
        self.cache = self.fileset[self.fileset_idx]
        self.fileset_idx += 1
        self.cache_idx = self.fileset_sub_idx

    def __call__(self, size=512):
        while len(self.cache) < self.cache_idx + size:
            if self.fileset_idx >= len(self.fileset):
                self.fileset_idx = 0
            self.fileset_sub_idx = self.cache_idx - len(self.cache)
            self.cache = self.cache[self.cache_idx:] + self.fileset[self.fileset_idx]
            self.cache_idx = 0
            self.fileset_idx += 1
        retn = self.cache[self.cache_idx:self.cache_idx + size]
        self.cache_idx += size
        self.fileset_sub_idx += size
        return retn

    def __repr__(self):
        return f"Hcorpus(r'{self.fileset.root}', fileset_idx={self.fileset_idx-1}, fileset_sub_idx={self.fileset_sub_idx})"

from tokenizers import Tokenizer
tokenizer = Tokenizer.from_file("HelloBPE.tokenizer.json")

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x * self.weight

class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim: int, max_seq_len: int, device=device, theta: float = 10000.0):
        super().__init__()
        self.head_dim = head_dim
        self.set_max_seq_len(max_seq_len, device, theta)

    def set_max_seq_len(self, max_seq_len: int, device=device, theta: float = 10000.0):
        self.max_seq_len = max_seq_len
        freqs = 1.0 / (theta ** (torch.arange(0, self.head_dim, 2).float().to(device) / self.head_dim))
        t = torch.arange(max_seq_len, device=device)  # type: ignore
        freqs = torch.outer(t, freqs).float()  # 外积
        self.freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # 复数，模 1，角度 freqs
        self.freqs_cis.requires_grad = False  # filter(lambda p : p.requires_grad, model.parameters())

    def rotary_emb(self, x):
        x_ = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
        x_out = torch.view_as_real(x_ * self.local_freqs_cis).flatten(3)
        return x_out.type_as(x)

    def forward(self, start_pos: int, seqlen: int):
        self.local_freqs_cis = self.freqs_cis[start_pos: start_pos + seqlen].view(1, seqlen, 1, -1)  # cacheKV 相关，可忽略
        self.local_freqs_cis.requires_grad = False
        return self.rotary_emb

class Attention(nn.Module):
    def __init__(self, hidden_size, n_heads, cacheKV, max_batch_size, max_seq_len, device=device):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = hidden_size // n_heads
        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)

    def forward(self, hidden_states, rotary_emb, start_pos=0, mask=None, is_causal=True):
        bsz, seqlen, hidden_size = hidden_states.shape

        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)

        q = q.view(bsz, seqlen, self.n_heads, self.head_dim)
        k = k.view(bsz, seqlen, self.n_heads, self.head_dim)
        v = v.view(bsz, seqlen, self.n_heads, self.head_dim)

        q = rotary_emb(q)
        k = rotary_emb(k)

        q = q.transpose(1, 2)  # (bs, n_heads, seqlen, head_dim)
        k = k.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        v = v.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)

        output = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=is_causal)

        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, hidden_size)
        return self.o_proj(output)

class MLP(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        intermediate_size = int(2 * hidden_size)
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x):
        gate = F.silu(self.gate_proj(x))
        intermediate_states = self.up_proj(x)
        return self.down_proj(gate * intermediate_states)

class Decoder(nn.Module):
    def __init__(self, hidden_size, n_heads, cacheKV, max_batch_size, max_seq_len):
        super().__init__()
        self.ln1 = RMSNorm(hidden_size)
        self.attn = Attention(hidden_size, n_heads, cacheKV, max_batch_size, max_seq_len)
        self.ln2 = RMSNorm(hidden_size)
        self.mlp = MLP(hidden_size)

    def forward(self, x, rotary_emb, start_pos, mask=None, is_causal=True):
        x = x + self.attn(self.ln1(x), rotary_emb, start_pos, mask, is_causal)
        return x + self.mlp(self.ln2(x))

class HelloGPT(nn.Module):
    def __init__(self, vocab_size=32765, hidden_size=768, n_heads=12, max_seq_len=1024, n_layers=12, cacheKV=False, max_batch_size=1):
        super().__init__()
        # hidden_size > 8.33 * ln(vocab_size)
        self.tok_embeddings = nn.Embedding(vocab_size, hidden_size)
        self.rotary_emb = RotaryEmbedding(hidden_size // n_heads, max_seq_len * 2)
        self.rotary_emb.requires_grad = False
        self.layers = nn.ModuleList()
        for layer_id in range(n_layers):
            self.layers.append(Decoder(hidden_size, n_heads, cacheKV, max_batch_size, max_seq_len))
        self.norm = RMSNorm(hidden_size)
        self.ln2 = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(self, input_ids: torch.Tensor, start_pos=0, no_mask=True):
        _bsz, seqlen = input_ids.shape
        h = self.tok_embeddings(input_ids)

        # 预计算，减少每一层的重复计算
        rotary_emb = self.rotary_emb(start_pos, seqlen)
        for layer in self.layers:
            h = layer(h, rotary_emb, start_pos)

        h = self.norm(h)
        h = self.ln2(h)
        return h.float()

data = Hcorpus(r'D:\datasets\h-corpus')
def get_batch(size=512, bsz=8):
    x = []
    y = []
    for i in range(bsz):
        tmp = data(size+1)
        x.append(tmp[:size])
        y.append(tmp[1:])
    return torch.tensor(x).to(device), torch.tensor(y).to(device)

model = HelloGPT(n_layers=8, max_seq_len=768)
model.to(device)

## 初始化训练器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数
optimizer = torch.optim.Adam(train_parameters, lr=6e-4)  # Adam 优化器
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=2)  # 余弦退火学习率
torch.manual_seed(1337)  # 魔术随机种子

total_loss = 0
print_iter = 20
for epoch in range(1, 100001):
    optimizer.zero_grad(set_to_none=True)  # 清空梯度，节省显存
    x, y = get_batch(size=384, bsz=4)  # x 是训练语料 y 是 x 移动了一位，当做预测目标
    y_ = model(x)  # 通过 x 预测的 y
    loss = criterion(y_.view(-1, 32765), y.view(-1))  # 计算损失
    loss.backward()  # 反向传播梯度
    torch.nn.utils.clip_grad_norm_(train_parameters, 0.5)  # 梯度裁剪，减轻过拟合
    optimizer.step()  # 通过梯度优化训练参数
    scheduler.step()  # 计算下一步的学习率
    total_loss += loss  # 累计损失

    if epoch % print_iter == 0:
        print(data)
        print(f'epoch: {epoch}  lr: {scheduler.get_last_lr()[0]:.4e} loss: {total_loss / print_iter:.4e}')
        total_loss = 0

with open('tmp_training.pkl', 'rb') as file:
    epoch = pickle.load(file)  # 读取 epoch 位置
    tmp_fileset_idx = pickle.load(file)  # 读取 data 位置
    tmp_fileset_sub_idx = pickle.load(file)
# 恢复数据位置
data = Hcorpus(r'D:\datasets\h-corpus', fileset_idx=tmp_fileset_idx-1, fileset_sub_idx=tmp_fileset_sub_idx)
model = torch.load(f'tmp_model_{epoch}.pth')  # 恢复模型
print(f'start from epoch: {epoch}   data: {data}')

save_iter = 5000
for epoch in range(1, 100001):
    pass
    if epoch % save_iter == 0:
        optimizer.zero_grad(set_to_none=True)  # 清空梯度，节省显存
        with open('tmp_training.pkl', 'wb') as file:
            pickle.dump(epoch, file)  # 保存 epoch 位置
            pickle.dump(data.fileset_idx, file)  # 保存 data 位置
            pickle.dump(data.fileset_sub_idx, file)
        torch.save(model, f'tmp_model_{epoch}.pth')  # 保存模型
        print(f'save to tmp_model_{epoch}.pth')

writer = SummaryWriter('logs')  # tensorboard --logdir logs
for epoch in range(1, 100001):
    pass
    writer.add_scalar('lr', scheduler.get_last_lr()[0], epoch)
    writer.add_scalar('loss', loss, epoch)
    if epoch % print_iter == 0:
        pass
        writer.add_scalar('total_loss', total_loss / print_iter, epoch)
writer.close()