/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
LOAD 'age';
SET search_path=ag_catalog;
SELECT create_graph('graph');
NOTICE:  graph "graph" has been created
 create_graph 
--------------
 
(1 row)

-- Should error out
SELECT * FROM cypher('graph', $$ RETURN cosine_distance("[1,2,3]", "[1,2,3]") $$) AS (n agtype);
ERROR:  function cosine_distance does not exist
LINE 1: SELECT * FROM cypher('graph', $$ RETURN cosine_distance("[1,...
                                       ^
HINT:  If the function is from an external extension, make sure the extension is installed and the function is in the search path.
-- Create the extension in the public schema
CREATE EXTENSION vector SCHEMA public;
-- Should error out
SELECT * FROM cypher('graph', $$ RETURN cosine_distance("[1,2,3]", "[1,2,3]") $$) AS (n agtype);
ERROR:  function cosine_distance does not exist
LINE 1: SELECT * FROM cypher('graph', $$ RETURN cosine_distance("[1,...
                                       ^
HINT:  If the function is from an external extension, make sure the extension is installed and the function is in the search path.
-- Should work
SET search_path=ag_catalog, public;
SELECT create_graph('graph');
ERROR:  graph "graph" already exists
SELECT * FROM cypher('graph', $$ RETURN "[1.22,2.22,3.33]"::vector $$) AS (n vector);
        n         
------------------
 [1.22,2.22,3.33]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN "[1.22,2.22,3.33]"::vector $$) AS (n halfvec);
                n                
---------------------------------
 [1.2197266,2.2207031,3.3300781]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN "[1.22,2.22,3.33]"::vector $$) AS (n sparsevec);
            n             
--------------------------
 {1:1.22,2:2.22,3:3.33}/3
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector $$) AS (n vector);
        n         
------------------
 [1.22,2.22,3.33]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector $$) AS (n halfvec);
                n                
---------------------------------
 [1.2197266,2.2207031,3.3300781]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector $$) AS (n sparsevec);
            n             
--------------------------
 {1:1.22,2:2.22,3:3.33}/3
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector $$) AS (n vector(3));
        n         
------------------
 [1.22,2.22,3.33]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector $$) AS (n halfvec(3));
                n                
---------------------------------
 [1.2197266,2.2207031,3.3300781]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector $$) AS (n sparsevec(3));
            n             
--------------------------
 {1:1.22,2:2.22,3:3.33}/3
(1 row)

-- Should error out
SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector $$) AS (n vector(2));
ERROR:  expected 2 dimensions, not 3
SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector $$) AS (n halfvec(2));
ERROR:  expected 2 dimensions, not 3
SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector $$) AS (n sparsevec(2));
ERROR:  expected 2 dimensions, not 3
SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector(3) $$) AS (n vector(4));
ERROR:  expected 4 dimensions, not 3
SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector(3) $$) AS (n halfvec(4));
ERROR:  expected 4 dimensions, not 3
SELECT * FROM cypher('graph', $$ RETURN [1.22,2.22,3.33]::vector(3) $$) AS (n sparsevec(4));
ERROR:  expected 4 dimensions, not 3
--
-- Test functions
--
SELECT * FROM cypher('graph', $$ RETURN l2_distance("[1,2,3]", "[1,2,4]") $$) AS (n agtype);
  n  
-----
 1.0
(1 row)

SELECT * FROM cypher('graph', $$ RETURN inner_product("[1,2,3]", "[1,2,4]") $$) AS (n agtype);
  n   
------
 17.0
(1 row)

SELECT * FROM cypher('graph', $$ RETURN cosine_distance("[1,2,3]", "[1,2,4]") $$) AS (n agtype);
          n          
---------------------
 0.00853986601633272
(1 row)

SELECT * FROM cypher('graph', $$ RETURN l1_distance("[1,2,3]", "[1,2,4]") $$) AS (n agtype);
  n  
-----
 1.0
(1 row)

SELECT * FROM cypher('graph', $$ RETURN vector_dims("[1,2,3]") $$) AS (n agtype);
 n 
---
 3
(1 row)

SELECT * FROM cypher('graph', $$ RETURN vector_norm("[1,2,3]") $$) AS (n agtype);
         n          
--------------------
 3.7416573867739413
(1 row)

SELECT * FROM cypher('graph', $$ RETURN l2_normalize("[1,2,3]") $$) AS (n vector);
                 n                 
-----------------------------------
 [0.26726124,0.5345225,0.80178374]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN l2_normalize("[1,2,3]")::text $$) AS (n agtype);
                  n                  
-------------------------------------
 [0.26726124, 0.5345225, 0.80178374]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN subvector("[1,2,3,4,5,6]", 2, 4) $$) AS (n vector);
     n     
-----------
 [2,3,4,5]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN subvector("[1,2,3,4,5,6]", 2, 4)::text $$) AS (n agtype);
      n       
--------------
 [2, 3, 4, 5]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN binary_quantize("[1,2,4]") $$) AS (n bit(3));
  n  
-----
 111
(1 row)

--
-- Test operators
--
SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector + [1,2,4]::vector $$) AS (n vector);
    n    
---------
 [2,4,7]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector - [1,2,4]::vector $$) AS (n vector);
    n     
----------
 [0,0,-1]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector * [1,2,4]::vector $$) AS (n vector);
    n     
----------
 [1,4,12]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector || [1,2,4]::vector $$) AS (n vector);
       n       
---------------
 [1,2,3,1,2,4]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector <#> [1,2,4]::vector $$) AS (n agtype);
   n   
-------
 -17.0
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector <=> [1,2,4]::vector $$) AS (n agtype);
          n          
---------------------
 0.00853986601633272
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector <+> [1,2,4]::vector $$) AS (n agtype);
  n  
-----
 1.0
(1 row)

--
-- Due to issues with pattern matching syntax, '-' is not allowed
-- as an operator character, so we have to use the OPERATOR syntax.
--
SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (`<->`) [1,2,4]::vector $$) AS (n agtype);
  n  
-----
 1.0
(1 row)

-- Using OPERATOR () syntax
SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (+) [1,2,4]::vector $$) AS (n vector);
    n    
---------
 [2,4,7]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (-) [1,2,4]::vector $$) AS (n vector);
    n     
----------
 [0,0,-1]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (*) [1,2,4]::vector $$) AS (n vector);
    n     
----------
 [1,4,12]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (||) [1,2,4]::vector $$) AS (n vector);
       n       
---------------
 [1,2,3,1,2,4]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (`<->`) [1,2,4]::vector $$) AS (n agtype);
  n  
-----
 1.0
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (<#>) [1,2,4]::vector $$) AS (n agtype);
   n   
-------
 -17.0
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (<=>) [1,2,4]::vector $$) AS (n agtype);
          n          
---------------------
 0.00853986601633272
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (<+>) [1,2,4]::vector $$) AS (n agtype);
  n  
-----
 1.0
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (public.+) [1,2,4]::vector $$) AS (n vector);
    n    
---------
 [2,4,7]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (public.-) [1,2,4]::vector $$) AS (n vector);
    n     
----------
 [0,0,-1]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (public.*) [1,2,4]::vector $$) AS (n vector);
    n     
----------
 [1,4,12]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (public.||) [1,2,4]::vector $$) AS (n vector);
       n       
---------------
 [1,2,3,1,2,4]
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (public.`<->`) [1,2,4]::vector $$) AS (n agtype);
  n  
-----
 1.0
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (public.<#>) [1,2,4]::vector $$) AS (n agtype);
   n   
-------
 -17.0
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (public.<=>) [1,2,4]::vector $$) AS (n agtype);
          n          
---------------------
 0.00853986601633272
(1 row)

SELECT * FROM cypher('graph', $$ RETURN [1,2,3]::vector OPERATOR (public.<+>) [1,2,4]::vector $$) AS (n agtype);
  n  
-----
 1.0
(1 row)

--
-- An example usage
--
SELECT * FROM cypher('graph', $$ 
    CREATE (:Movie {title: "The Matrix", year: 1999, genre: "Action", plot: "A computer hacker learns about the true nature of reality and joins a rebellion to free humanity from a simulated world controlled by machines.", embedding: [-0.07594558, 0.04081754, 0.29592122, -0.11921061]}),
           (:Movie {title: "The Matrix Reloaded", year: 2003, genre: "Action", plot: "The rebels continue their fight against the machines, uncovering deeper truths about the Matrix and the nature of their mission.", embedding: [0.30228977, -0.22839354, 0.35070436, 0.01262819]}),
           (:Movie {title: "The Matrix Revolutions", year: 2003, genre: "Action", plot: "The final battle between humans and machines reaches its climax as the fate of both worlds hangs in the balance.", embedding: [ 0.12240622, -0.29752459, 0.22620453, 0.24454723]}),
           (:Movie {title: "The Matrix Resurrections", year: 2021, genre: "Action", plot: "Neo returns to a new version of the Matrix and must once again fight to save the people from the control of the machines.", embedding: [ 0.34717246, -0.13820869, 0.29214213, 0.08090488]}),
           (:Movie {title: "Inception", year: 2010, genre: "Sci-Fi", plot: "A skilled thief is given a chance at redemption if he can successfully perform an inception: planting an idea into someone’s subconscious.", embedding: [ 0.03923657, 0.39284106, -0.20927092, -0.17770818]}),
           (:Movie {title: "Interstellar", year: 2014, genre: "Sci-Fi", plot: "A group of explorers travel through a wormhole in space in an attempt to ensure humanity’s survival.", embedding: [-0.29302418, -0.39615033, -0.23393948, -0.09601383]}),
           (:Movie {title: "Avatar", year: 2009, genre: "Sci-Fi", plot: "A paraplegic Marine is sent to the moon Pandora, where he becomes torn between following orders and protecting the world he feels is his home.", embedding: [-0.13663386, 0.00635589, -0.03038832, -0.08252723]}),
           (:Movie {title: "Blade Runner", year: 1982, genre: "Sci-Fi", plot: "A blade runner must pursue and terminate four replicants who have stolen a ship in space and returned to Earth.", embedding: [ 0.27215557, -0.1479577, -0.09972772, -0.08234394]}),
           (:Movie {title: "Blade Runner 2049", year: 2017, genre: "Sci-Fi", plot: "A new blade runner unearths a long-buried secret that has the potential to plunge what’s left of society into chaos.", embedding: [ 0.21560573, -0.07505179, -0.01331814, 0.13403069]}),
           (:Movie {title: "Minority Report", year: 2002, genre: "Sci-Fi", plot: "In a future where a special police unit can arrest murderers before they commit their crimes, a top officer is accused of a future murder.", embedding: [ 0.24008012, 0.44954908, -0.30905488, 0.15195407]}),
           (:Movie {title: "Total Recall", year: 1990, genre: "Sci-Fi", plot: "A construction worker discovers that his memories have been implanted and becomes embroiled in a conspiracy on Mars.", embedding: [-0.17471036, 0.14695261, -0.06272433, -0.21795064]}),
           (:Movie {title: "Elysium", year: 2013, genre: "Sci-Fi", plot: "In a future where the rich live on a luxurious space station while the rest of humanity lives in squalor, a man fights to bring equality.", embedding: [-0.33280967, 0.07733926, 0.11015328, 0.53382836]}),
           (:Movie {title: "Gattaca", year: 1997, genre: "Sci-Fi", plot: "In a future where genetic engineering determines social class, a man defies his fate to achieve his dreams.", embedding: [-0.21629286, 0.31114665, 0.08303899, 0.46199759]}),
           (:Movie {title: "The Fifth Element", year: 1997, genre: "Sci-Fi", plot: "In a futuristic world, a cab driver becomes the key to saving humanity from an impending cosmic threat.", embedding: [-0.11528205, -0.0208782, -0.0735215, 0.14327449]}),
           (:Movie {title: "The Terminator", year: 1984, genre: "Action", plot: "A cyborg assassin is sent back in time to kill the mother of the future resistance leader.", embedding: [ 0.33666933, 0.18040994, -0.01075103, -0.11117851]}),
           (:Movie {title: "Terminator 2: Judgment Day", year: 1991, genre: "Action", plot: "A reprogrammed Terminator is sent to protect the future leader of the human resistance from a more advanced Terminator.", embedding: [ 0.34698868, 0.06439331, 0.06232323, -0.19534876]}),
           (:Movie {title: "Jurassic Park", year: 1993, genre: "Adventure", plot: "Scientists clone dinosaurs to create a theme park, but things go awry when the creatures escape.", embedding: [ 0.01794725, -0.11434246, -0.46831815, -0.01049593]}),
           (:Movie {title: "The Avengers", year: 2012, genre: "Action", plot: "Superheroes assemble to face a global threat from an alien invasion led by Loki.", embedding: [ 0.00546514, -0.37005171, -0.42612838, 0.07968612]})
$$) AS (result agtype);
 result 
--------
(0 rows)

SELECT * FROM cypher('graph', $$ MATCH (m:Movie) RETURN m.title, (m.embedding)::vector $$) AS (title agtype, embedding vector);
            title             |                     embedding                     
------------------------------+---------------------------------------------------
 "The Matrix"                 | [-0.07594558,0.04081754,0.2959212,-0.11921061]
 "The Matrix Reloaded"        | [0.30228978,-0.22839354,0.35070437,0.01262819]
 "The Matrix Revolutions"     | [0.12240622,-0.2975246,0.22620453,0.24454723]
 "The Matrix Resurrections"   | [0.34717247,-0.13820869,0.29214212,0.08090488]
 "Inception"                  | [0.03923657,0.39284107,-0.20927092,-0.17770818]
 "Interstellar"               | [-0.29302418,-0.39615032,-0.23393948,-0.09601383]
 "Avatar"                     | [-0.13663386,0.00635589,-0.03038832,-0.08252723]
 "Blade Runner"               | [0.27215558,-0.1479577,-0.09972772,-0.08234394]
 "Blade Runner 2049"          | [0.21560574,-0.07505179,-0.01331814,0.13403068]
 "Minority Report"            | [0.24008012,0.44954908,-0.30905488,0.15195407]
 "Total Recall"               | [-0.17471036,0.14695261,-0.06272433,-0.21795064]
 "Elysium"                    | [-0.33280966,0.07733926,0.11015328,0.5338284]
 "Gattaca"                    | [-0.21629286,0.31114665,0.08303899,0.4619976]
 "The Fifth Element"          | [-0.11528205,-0.0208782,-0.0735215,0.14327449]
 "The Terminator"             | [0.33666933,0.18040994,-0.01075103,-0.11117851]
 "Terminator 2: Judgment Day" | [0.34698868,0.06439331,0.06232323,-0.19534875]
 "Jurassic Park"              | [0.01794725,-0.11434246,-0.46831816,-0.01049593]
 "The Avengers"               | [0.00546514,-0.3700517,-0.4261284,0.07968612]
(18 rows)

-- Check the dimension of the embedding
SELECT * FROM cypher('graph', $$ MATCH (m:Movie) RETURN m.title, vector_dims(m.embedding) $$) AS (title agtype, dimension int);
            title             | dimension 
------------------------------+-----------
 "The Matrix"                 |         4
 "The Matrix Reloaded"        |         4
 "The Matrix Revolutions"     |         4
 "The Matrix Resurrections"   |         4
 "Inception"                  |         4
 "Interstellar"               |         4
 "Avatar"                     |         4
 "Blade Runner"               |         4
 "Blade Runner 2049"          |         4
 "Minority Report"            |         4
 "Total Recall"               |         4
 "Elysium"                    |         4
 "Gattaca"                    |         4
 "The Fifth Element"          |         4
 "The Terminator"             |         4
 "Terminator 2: Judgment Day" |         4
 "Jurassic Park"              |         4
 "The Avengers"               |         4
(18 rows)

-- Get top 4 most similar movies to The Terminator using cosine distance
SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Terminator"})
                                 RETURN m.title ORDER BY cosine_distance(m.embedding, search.embedding)
                                 ASC LIMIT 4
$$) AS (title agtype);
            title             
------------------------------
 "The Terminator"
 "Terminator 2: Judgment Day"
 "Minority Report"
 "Blade Runner"
(4 rows)

SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Terminator"})
                                 RETURN m.title ORDER BY m.embedding::vector <=> search.embedding::vector
                                 ASC LIMIT 4
$$) AS (title agtype);
            title             
------------------------------
 "The Terminator"
 "Terminator 2: Judgment Day"
 "Minority Report"
 "Blade Runner"
(4 rows)

-- Get top 4 most similar movies to The Matrix using cosine distance
SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Matrix"})
                                 RETURN m.title ORDER BY cosine_distance(m.embedding, search.embedding)
                                 ASC LIMIT 4
$$) AS (title agtype);
           title            
----------------------------
 "The Matrix"
 "The Matrix Reloaded"
 "The Matrix Resurrections"
 "Total Recall"
(4 rows)

SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Matrix"})
                                 RETURN m.title ORDER BY m.embedding::vector <=> search.embedding::vector
                                 ASC LIMIT 4
$$) AS (title agtype);
           title            
----------------------------
 "The Matrix"
 "The Matrix Reloaded"
 "The Matrix Resurrections"
 "Total Recall"
(4 rows)

-- l2 norm of the embedding
SELECT * FROM cypher('graph', $$ MATCH (m:Movie) set m.embedding=l2_normalize(m.embedding)::agtype return m.title, m.embedding $$) AS (title agtype, embedding agtype);
            title             |                      embedding                      
------------------------------+-----------------------------------------------------
 "The Matrix"                 | [-0.22980669, 0.12351139, 0.89543957, -0.36072403]
 "The Matrix Reloaded"        | [0.58534974, -0.44225806, 0.6790991, 0.024453051]
 "The Matrix Revolutions"     | [0.26431033, -0.6424414, 0.4884408, 0.528048]
 "The Matrix Resurrections"   | [0.72151977, -0.28723562, 0.60715157, 0.16814256]
 "Inception"                  | [0.08159459, 0.81693435, -0.43519026, -0.3695538]
 "Interstellar"               | [-0.5290723, -0.71527255, -0.4223914, -0.17335857]
 "Avatar"                     | [-0.84023285, 0.039085682, -0.18687363, -0.507503]
 "Blade Runner"               | [0.81074023, -0.44075987, -0.29708475, -0.2452992]
 "Blade Runner 2049"          | [0.8134027, -0.28314334, -0.05024454, 0.50564945]
 "Minority Report"            | [0.39031598, 0.7308651, -0.5024533, 0.24704295]
 "Total Recall"               | [-0.54291505, 0.4566574, -0.19491677, -0.67728484]
 "Elysium"                    | [-0.517338, 0.12022049, 0.17122844, 0.82981277]
 "Gattaca"                    | [-0.35853538, 0.51576865, 0.13764863, 0.765825]
 "The Fifth Element"          | [-0.5788842, -0.10483904, -0.36918527, 0.7194471]
 "The Terminator"             | [0.84599304, 0.45333964, -0.02701552, -0.27937278]
 "Terminator 2: Judgment Day" | [0.8501332, 0.15776564, 0.15269388, -0.4786106]
 "Jurassic Park"              | [0.037194606, -0.23696794, -0.9705615, -0.02175219]
 "The Avengers"               | [0.009587915, -0.6492101, -0.7475897, 0.13979948]
(18 rows)

-- Get top 4 most similar movies to The Terminator using l2 distance
SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Terminator"})
                                 RETURN m.title ORDER BY l2_distance(m.embedding, search.embedding) ASC LIMIT 4
$$) AS (title agtype);
            title             
------------------------------
 "The Terminator"
 "Terminator 2: Judgment Day"
 "Minority Report"
 "Blade Runner"
(4 rows)

SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Terminator"})
                                 RETURN m.title ORDER BY m.embedding::vector OPERATOR (`<->`) search.embedding::vector
                                 ASC LIMIT 4
$$) AS (title agtype);
            title             
------------------------------
 "The Terminator"
 "Terminator 2: Judgment Day"
 "Minority Report"
 "Blade Runner"
(4 rows)

-- Get top 4 most similar movies to The Matrix using l2 distance
SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Matrix"})
                                 RETURN m.title ORDER BY l2_distance(m.embedding, search.embedding) ASC LIMIT 4
$$) AS (title agtype);
           title            
----------------------------
 "The Matrix"
 "The Matrix Reloaded"
 "The Matrix Resurrections"
 "Total Recall"
(4 rows)

SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Matrix"})
                                 RETURN m.title ORDER BY m.embedding::vector OPERATOR (`<->`) search.embedding::vector
                                 ASC LIMIT 4
$$) AS (title agtype);
           title            
----------------------------
 "The Matrix"
 "The Matrix Reloaded"
 "The Matrix Resurrections"
 "Total Recall"
(4 rows)

--
-- Test vector index
--
-- This function will be used to check if index scan
-- is used successfully. We cannot simply have EXPLAIN
-- in the upcoming queries because it produces some 
-- hardcoded oids in sort node, which may change in
-- future and break the tests.
CREATE OR REPLACE FUNCTION plan_has_index_scan(sql text)
RETURNS boolean
LANGUAGE plpgsql AS
$$               
DECLARE                                                                                   
    plan_lines text[];                                                                             
    plan_text text;
BEGIN                
    EXECUTE format('EXPLAIN (FORMAT JSON, COSTS OFF) %s', sql) INTO plan_text;

    -- Return true if 'Index Scan' appears anywhere
    RETURN position('"Index Scan"' in plan_text) > 0;
END;
$$;
SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Matrix"})
                                 RETURN m.title ORDER BY m.embedding::vector(4) <=> search.embedding::vector(4)
                                 ASC LIMIT 4
$$) AS (title agtype);
           title            
----------------------------
 "The Matrix"
 "The Matrix Reloaded"
 "The Matrix Resurrections"
 "Total Recall"
(4 rows)

-- The index expression below matches the expression
-- seen in the EXPLAIN plan of above query
DO $$
DECLARE
    graph_oid oid;
BEGIN
    SELECT graphid INTO graph_oid
    FROM ag_catalog.ag_graph
    WHERE name = 'graph';

    EXECUTE format($f$
        CREATE INDEX movie_vector_idx ON graph."Movie"
        USING hnsw (((
          agtype_access_operator(
            VARIADIC ARRAY[
              _agtype_build_vertex(id, _label_name(%L::oid, id), properties),
              '"embedding"'::agtype
            ]
          )::text
        )::vector(4)) vector_cosine_ops);
    $f$, graph_oid);
END;
$$;
-- Disable seqscan just to test the index
SET enable_seqscan = off;
SELECT plan_has_index_scan($f$
    SELECT * FROM cypher('graph', $$
        MATCH (m:Movie)
        RETURN m.title
        ORDER BY m.embedding::vector(4) <=> [-0.07594558, 0.04081754, 0.29592122, -0.11921061]::vector(4)
        ASC LIMIT 4
    $$) AS (title agtype);
$f$);
 plan_has_index_scan 
---------------------
 t
(1 row)

SELECT * FROM cypher('graph', $$ MATCH (m:Movie)
                                 RETURN m.title
                                 ORDER BY m.embedding::vector(4) <=> [-0.07594558, 0.04081754, 0.29592122, -0.11921061]::vector(4)
                                 ASC LIMIT 4
$$) AS (title agtype);
           title            
----------------------------
 "The Matrix"
 "The Matrix Reloaded"
 "The Matrix Resurrections"
 "Total Recall"
(4 rows)

DROP INDEX graph.movie_vector_idx;
SET enable_seqscan = on;
-- Test a direct implicit cast
CREATE CAST (agtype AS vector)
    WITH INOUT AS implicit;
SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Matrix"})
                                 RETURN m.title ORDER BY m.embedding <=> search.embedding
                                 ASC LIMIT 4
$$) AS (title agtype);
           title            
----------------------------
 "The Matrix"
 "The Matrix Reloaded"
 "The Matrix Resurrections"
 "Total Recall"
(4 rows)

SELECT * FROM cypher('graph', $$ MATCH (m:Movie), (search:Movie {title: "The Matrix"})
                                 RETURN m.title ORDER BY m.embedding OPERATOR (`<->`) search.embedding
                                 ASC LIMIT 4
$$) AS (title agtype);
           title            
----------------------------
 "The Matrix"
 "The Matrix Reloaded"
 "The Matrix Resurrections"
 "Total Recall"
(4 rows)

DO $$
DECLARE
    graph_oid oid;
BEGIN
    SELECT graphid INTO graph_oid
    FROM ag_catalog.ag_graph
    WHERE name = 'graph';

    EXECUTE format($f$
        CREATE INDEX movie_vector_idx ON graph."Movie"
        USING hnsw ((
          agtype_access_operator(
            VARIADIC ARRAY[
              _agtype_build_vertex(id, _label_name(%L::oid, id), properties),
              '"embedding"'::agtype
            ]
        )::vector(4)) vector_cosine_ops);
    $f$, graph_oid);
END;
$$;
-- Disable seqscan just to test the index
SET enable_seqscan = off;
SELECT plan_has_index_scan($f$
    SELECT * FROM cypher('graph', $$
        MATCH (m:Movie)
        RETURN m.title
        ORDER BY m.embedding::vector(4) <=> [-0.07594558, 0.04081754, 0.29592122, -0.11921061]::vector(4)
        ASC LIMIT 4
    $$) AS (title agtype);
$f$);
 plan_has_index_scan 
---------------------
 t
(1 row)

SELECT * FROM cypher('graph', $$ MATCH (m:Movie)
                                 RETURN m.title
                                 ORDER BY m.embedding::vector(4) <=> [-0.07594558, 0.04081754, 0.29592122, -0.11921061]::vector(4)
                                 ASC LIMIT 4
$$) AS (title agtype);
           title            
----------------------------
 "The Matrix"
 "The Matrix Reloaded"
 "The Matrix Resurrections"
 "Total Recall"
(4 rows)

SET enable_seqscan = on;
--
-- Clean up
--
DROP FUNCTION plan_has_index_scan(text);
DROP CAST (agtype AS vector);
SELECT drop_graph('graph', true);
NOTICE:  drop cascades to 3 other objects
DETAIL:  drop cascades to table graph._ag_label_vertex
drop cascades to table graph._ag_label_edge
drop cascades to table graph."Movie"
NOTICE:  graph "graph" has been dropped
 drop_graph 
------------
 
(1 row)

DROP EXTENSION vector CASCADE;