import torch
import torchvision
from torchvision.transforms import (
    Compose, 
    Resize, 
    CenterCrop, 
    ToTensor, 
    Normalize
)
preprocess = Compose([
    Resize(256),
    CenterCrop(224),
    ToTensor(),
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
DATA_DIRECTORY = "/Users/shangwenkai/Desktop/vector/CIFAR"
datasets = {
    "CIFAR100": torchvision.datasets.CIFAR100(DATA_DIRECTORY, transform=preprocess, download=True)
}

# 准备数据。
BATCH_SIZE = 100
dataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE)
# 下载模型。
# 预训练的SqueezeNet1_1模型，并将其设置为评估模式。在评估模式下，模型不进行梯度计算，这对于特征提取是必要的，因为不需要反向传播
model = torchvision.models.squeezenet1_1(pretrained=True).eval()
# 提取特征向量，写入features_file_path。
features_file_path = "/Users/shangwenkai/Desktop/vector/features"
feature_file = open(features_file_path, 'w')
img_id = 0
for batch_number, batch in enumerate(dataloader): #每次处理100
    with torch.no_grad(): #禁用梯度
        batch_imgs = batch[0]  # 0: images
        batch_labels = batch[1]  # 1: labels
        vector_values = model(batch_imgs).tolist()
        #对batch_imgs进行前向传播，得到每个图像的特征向量，然后将这些特征向量从Tensor转换为Python列表。
        #SqueezeNet模型的输出是一个形状为(batch_size, output_dim)的二维Tensor
        for i in range(len(vector_values)):
            img_label = dataset.classes[batch_labels[i].item()]
            # print(img_label)
            feature_file.write(str(img_id) + "|" + img_label + "|")
            
            vector_value = vector_values[i]
            assert len(vector_value) == 1000
            for j in range(len(vector_value)):
                if j == 0:
                    feature_file.write("{")
                    feature_file.write(str(vector_value[j]) + ",")
                elif j == len(vector_value) - 1:
                    feature_file.write(str(vector_value[j]))
                    feature_file.write("}")
                else:
                    feature_file.write(str(vector_value[j]) + ",")
            feature_file.write("\n")
            
            img_id = img_id + 1
        print("finished extract feature vector for batch: ", batch_number)
feature_file.close()

import os
import psycopg2cffi
connection = psycopg2cffi.connect(
    host=os.environ.get("PGHOST", "pgm-bp16r2431088910yno.pg.rds.aliyuncs.com"),
    port=os.environ.get("PGPORT", "5432"),
    database=os.environ.get("PGDATABASE", "face_ai"),
    user=os.environ.get("PGUSER", "zhongchi"),
    password=os.environ.get("PGPASSWORD", "Ssaa1776")
)
cursor = connection.cursor()
drop_table_sql = """
drop TABLE IF EXISTS public.image_search ;
"""
# 用于创建表的SQL语句。
create_table_sql = """
CREATE TABLE IF NOT EXISTS public.image_search (
    id INTEGER NOT NULL,
    class TEXT,
    image_vector vector(1000),
    PRIMARY KEY(id)
) ;
"""
# 修改向量列的存储格式为PLAIN。
alter_vector_storage_sql = """
ALTER TABLE public.image_search ALTER COLUMN image_vector SET STORAGE PLAIN;
"""
# 执行上述SQL语句。
#cursor.execute(drop_table_sql)
cursor.execute(create_table_sql)
cursor.execute(alter_vector_storage_sql)
connection.commit()
print("rds pg connect over!")

import io
# 定义一个生成器函数，逐行处理文件中的数据。
def process_file(file_path):
    with open(file_path, 'r') as file:
        for line in file:
            modified_line = line.replace('{', '[').replace('}', ']')
            yield modified_line
# 导入数据的SQL。
copy_command = """
COPY public.image_search (id, class, image_vector)
FROM STDIN WITH (DELIMITER '|');
"""
# 图片特征向量文件。
features_file_path = "/Users/shangwenkai/Desktop/vector/features"
# 执行COPY命令。
modified_lines = io.StringIO(''.join(list(process_file(features_file_path))))
print("uploadeing")
cursor.copy_expert(copy_command, modified_lines)
connection.commit()
print("uploaded")

#CREATE INDEX ON vtest USING ivfflat(v vector_cosine_ops) WITH(lists = 100); 
#CREATE INDEX ON vtest USING ivfflat(v vector_cosine_ops) WITH(lists = 200); 
#CREATE INDEX ON vtest USING ivfflat(v vector_cosine_ops) WITH(lists = 500); 
#CREATE INDEX ON vtest USING hnsw (v vector_cosine_ops)
#CREATE INDEX ON vtest USING hnsw (v vector_cosine_ops) with(m = 32, ef_construction = 250); 
#CREATE INDEX ON vtest USING hnsw (v vector_cosine_ops) with(m = 32, ef_construction = 500); 
create_indexes_sql = """
CREATE INDEX ON image_search USING ivfflat(image_vector vector_l2_ops) WITH(lists = 100);
"""
print("creating index!")
cursor.execute(create_indexes_sql)
connection.commit()
print("rds pg index over!")

import os
import psycopg2cffi
def show_images_from_full_dataset(dset, num_rows, num_cols, indices):        
    im_arrays = np.take(dset.data, indices, axis=0)
    labels = map(dset.classes.__getitem__, np.take(dset.targets, indices))
    indices = np.array(indices)  
    fig = plt.figure(figsize=(18, 18))
    grid = ImageGrid(
        fig, 
        111,
        nrows_ncols=(num_rows, num_cols),
        axes_pad=0.3)
    for i, (ax, im_array, label, idx) in enumerate(zip(grid, im_arrays, labels, indices)):
        ax.imshow(im_array)
        # 在标题中显示编号和标签
        ax.set_title(f"{idx}: {label}")
        ax.axis("off")
def query_analyticdb(collection_name, vector_name, query_embedding, top_k=50):
# 创建查询SQL，返回与查询向量最相近的图片，同时计算与查询向量的相似度。
    query_sql = f"""
    SELECT id, class, ({vector_name} <=> '{query_embedding}') AS cosine_distance
    FROM {collection_name}
    ORDER BY cosine_distance
    LIMIT {top_k};
    """
    #print(query_sql)
# 执行查询。
    connection = psycopg2cffi.connect(
        host=os.environ.get("PGHOST", "pgm-bp16r2431088910yno.pg.rds.aliyuncs.com"),
        port=os.environ.get("PGPORT", "5432"),
        database=os.environ.get("PGDATABASE", "face_ai"),
        user=os.environ.get("PGUSER", "zhongchi"),
        password=os.environ.get("PGPASSWORD", "Ssaa1776")
    )
    
    cursor = connection.cursor()
    connection.autocommit = False
    cursor.execute("SET LOCAL ivfflat.probes = 100;")
    
    cursor.execute(query_sql)
    results = cursor.fetchall()
    connection.commit()
    return results
  
# 选择一条数据作为query。
def select_feature(file_path, expect_id):
    with open(file_path, 'r') as file:
        for line in file:
            datas = line.split('|')
            if datas[0] == str(expect_id):
                vec = '[' + datas[2][1:-2] + ']'
                return vec
    raise ValueError(f"没有对应id= {expect_id}的数据")
file_path = "/Users/shangwenkai/Desktop/vector/features"
# 选取id为4999的图片。
num=1
query_vector = select_feature(file_path, num)
# 查看这张图片。
show_images_from_full_dataset(dataset, 1, 1, [num])
#print(query_vector)
# 执行查询。
results = query_analyticdb("image_search", "image_vector", query_vector)
# 将查询结果对应的图片显示出来
# 获取上一步返回结果中的图片id。
indices = []
for item in results:
    indices.append(item[0])
print(indices)
# 显示图片。
show_images_from_full_dataset(dataset, 5, 10, indices)

CREATE TABLE vtest(id BIGINT, v VECTOR(1536));

WITH tmp AS (
    SELECT random_array(1536)::VECTOR(1536) AS vec
)
SELECT
  id,
  1 - (v <=> (SELECT vec FROM tmp)) AS cosine_similarity
FROM
  vtest
ORDER BY
  v <=> (SELECT vec FROM tmp)
LIMIT FLOOR(RANDOM() * 50);

QUERY PLAN
Limit  (cost=743.55..885.11 rows=10000 width=24) (actual time=13.342..13.788 rows=38 loops=1)
  
  CTE tmp
    ->  Result  (cost=0.00..0.01 rows=1 width=32) (actual time=0.340..0.340 rows=1 loops=1)
  
  InitPlan 2 (returns $1) 
    ->  CTE Scan on tmp  (cost=0.00..0.02 rows=1 width=32) (actual time=0.002..0.002 rows=1 loops=1)
  
  InitPlan 3 (returns $2) 
    ->  CTE Scan on tmp tmp_1  (cost=0.00..0.02 rows=1 width=32) (actual time=0.343..0.344 rows=1 loops=1)
  
  ->  Index Scan using vtest_v_idx on vtest  (cost=743.50..2159.10 rows=100000 width=24) (actual time=13.302..13.744 rows=38 loops=1)
        Order By: (v <=> $2)
  
Planning Time: 0.413 ms
Execution Time: 13.919 ms

-- 可以
ORDER BY embedding <=> '[3,1,2]' LIMIT 5;
-- 不可以
ORDER BY 1 - (embedding <=> '[3,1,2]') DESC LIMIT 5;

SET max_parallel_maintenance_workers = 4;
SET maintenance_work_mem = '7GB';

SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;

CREATE TABLE items (id bigserial PRIMARY KEY, embedding halfvec(3));
CREATE INDEX ON items USING hnsw ((embedding::halfvec(3)) halfvec_l2_ops);
SELECT * FROM items ORDER BY embedding::halfvec(3) <-> '[1,2,3]' LIMIT 5;

CREATE TABLE items (id bigserial PRIMARY KEY, embedding sparsevec(5));
INSERT INTO items (embedding) VALUES ('{1:1,3:2,5:3}/5'), ('{1:4,3:5,5:6}/5');
SELECT * FROM items ORDER BY embedding <-> '{1:3,3:1,5:2}/5' LIMIT 5;

CREATE INDEX ON items USING hnsw ((subvector(embedding, 1, 3)::vector(3)) vector_cosine_ops);
SELECT * FROM items ORDER BY subvector(embedding, 1, 3)::vector(3) <=> subvector('[1,2,3,4,5]'::vector, 1, 3) LIMIT 5;

SELECT * FROM (
    SELECT * FROM items ORDER BY subvector(embedding, 1, 3)::vector(3) <=> subvector('[1,2,3,4,5]'::vector, 1, 3) LIMIT 20
) ORDER BY embedding <=> '[1,2,3,4,5]' LIMIT 5;

SELECT
  id,
  1 - (v <=> '目标向量') AS cosine_similarity
FROM
  vtest
ORDER BY
  v <=> '目标向量'
LIMIT 50;

Limit  (cost=89.70..90.98 rows=50 width=24) (actual time=1.720..2.044 rows=40 loops=1)
  ->  Index Scan using vtest_v_idx3 on vtest  (cost=89.70..2649.30 rows=100000 width=24) (actual time=1.719..2.039 rows=40 loops=1)
        Order By: (v <=> '目标向量'::vector)
Planning Time: 0.080 ms
Execution Time: 2.081 ms

SELECT
  t.id as id,
  (1.0 - t.score) as score
FROM
(
SELECT
  id,
  (v <=> '目标向量') 
  AS score
FROM
  vtest
ORDER BY
  score
LIMIT 50
) t;

Subquery Scan on t  (cost=89.70..90.85 rows=50 width=16) (actual time=1.616..1.845 rows=40 loops=1)
  ->  Limit  (cost=89.70..90.73 rows=50 width=16) (actual time=1.615..1.839 rows=40 loops=1)
        ->  Index Scan using vtest_v_idx3 on vtest  (cost=89.70..2149.30 rows=100000 width=16) (actual time=1.614..1.834 rows=40 loops=1)
              Order By: (v <=> '目标向量'::vector)
Planning Time: 0.088 ms
Execution Time: 1.880 ms

SELECT
  id,
  1 - (v <=> '目标向量') AS cosine_similarity
FROM
  vtest
WHERE v <=> '目标向量'
  <=0.3
ORDER BY
  v <=> '目标向量'
LIMIT 50
;

Limit  (cost=89.70..93.54 rows=50 width=24) (actual time=1.699..2.152 rows=40 loops=1)
  ->  Index Scan using vtest_v_idx3 on vtest  (cost=89.70..2649.30 rows=33333 width=24) (actual time=1.698..2.147 rows=40 loops=1)
        Order By: (v <=> '目标向量'::vector)
        Filter: ((v <=> '目标向量'::vector) <= '0.3'::double precision)
Planning Time: 0.107 ms
Execution Time: 2.187 ms

SELECT
  t.id as id,
  (1.0 - t.score) as score
FROM
(
SELECT
  id,
  (v <=> '目标向量') 
  AS score
FROM
  vtest
ORDER BY
  score
LIMIT 50
) t
WHERE score<=0.3;

Subquery Scan on t  (cost=89.70..91.40 rows=17 width=16) (actual time=1.767..2.007 rows=40 loops=1)
  Filter: (t.score <= '0.3'::double precision)
  ->  Limit  (cost=89.70..90.73 rows=50 width=16) (actual time=1.766..1.999 rows=40 loops=1)
        ->  Index Scan using vtest_v_idx3 on vtest  (cost=89.70..2149.30 rows=100000 width=16) (actual time=1.765..1.994 rows=40 loops=1)
              Order By: (v <=> '目标向量'::vector)
Planning Time: 0.433 ms
Execution Time: 2.078 ms

SELECT id, content FROM items, plainto_tsquery('hello search') query
    WHERE textsearch @@ query ORDER BY ts_rank_cd(textsearch, query) DESC LIMIT 5;