import asyncio
import os
from typing import Optional, Dict, List, Any
from dataclasses import dataclass
import time
import torch
from torch import amp
import torch.nn as nn
import numpy as np
from collections import deque

from solana_logger import log_debug, log_info, log_error, SolanaLogger

logger = SolanaLogger()

class FeatureExtractor(nn.Module):
    def __init__(self, input_dim=512, hidden_dims=[256, 128]):
        super().__init__()
        layers = []
        prev_dim = input_dim
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.LeakyReLU(0.01),
                nn.Dropout(0.1)
            ])
            prev_dim = hidden_dim
        self.feature_extractor = nn.Sequential(*layers)
        self.projection_head = nn.Linear(prev_dim, 64)

    def forward(self, x):
        features = self.feature_extractor(x)
        return self.projection_head(features)

@dataclass
class VectorConfig:
    batch_size: int = int(os.getenv("VECTOR_BATCH_SIZE", "64"))
    vector_dim: int = int(os.getenv("VECTOR_DIMENSION", "512"))
    similarity_threshold: float = float(os.getenv("SIMILARITY_THRESHOLD", "0.85"))
    cache_size: int = int(os.getenv("VECTOR_CACHE_SIZE", "1000"))

@dataclass
class VectorMetrics:
    patterns_stored: int = 0
    patterns_retrieved: int = 0
    last_cleanup: float = time.time()

class EnhancedVectorStore:
    def __init__(self):
        self.config = VectorConfig()
        self.metrics = VectorMetrics()
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self._setup_model()
        self._init_storage()

    def _setup_model(self):
        self.feature_extractor = FeatureExtractor().to(self.device)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.backends.cudnn.benchmark = True
        self.scaler = amp.GradScaler()

    def _init_storage(self):
        self.patterns = {}
        self.cache = {}
        self.recent_patterns = deque(maxlen=self.config.cache_size)

    async def __aenter__(self):
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.cleanup()

    async def initialize(self):
        """Initialize resources"""
        try:
            log_info("Initializing vector store")
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            return True
        except Exception as e:
            log_error(f"Initialization failed: {e}")
            return False

    def _calculate_similarity(self, vec1: torch.Tensor, vec2: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            vec1_norm = torch.nn.functional.normalize(vec1, p=2, dim=1)
            vec2_norm = torch.nn.functional.normalize(vec2, p=2, dim=1)
            return torch.mm(vec1_norm, vec2_norm.t())

    async def store_pattern(self, pool_id: str, pattern_vector: torch.Tensor, metadata: Dict):
        try:
            pattern_vector = pattern_vector.to(device=self.device, dtype=torch.float16)
            with amp.autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu'):
                features = self.feature_extractor(pattern_vector)
            
            self.patterns[pool_id] = {
                'vector': features,
                'timestamp': time.time(),
                'metadata': metadata
            }
            self.cache[pool_id] = features
            self.recent_patterns.append(pool_id)
            self.metrics.patterns_stored += 1

            if len(self.patterns) > self.config.cache_size:
                await self._cleanup_old_patterns()

        except Exception as e:
            log_error(f"Pattern storage failed: {e}")

    async def search_similar_patterns(self, query_vector: torch.Tensor, k: int = 5) -> List[Dict]:
        try:
            if not self.patterns:
                return []

            query_vector = query_vector.to(device=self.device, dtype=torch.float16)
            with amp.autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu'):
                query_features = self.feature_extractor(query_vector)

            pattern_vectors = torch.stack([p['vector'] for p in self.patterns.values()], dim=0)
            similarities = self._calculate_similarity(query_features.unsqueeze(0), pattern_vectors)

            results = []
            for pool_id, sim in zip(self.patterns.keys(), similarities[0]):
                if sim >= self.config.similarity_threshold:
                    results.append({
                        'pool_id': pool_id,
                        'similarity': float(sim),
                        'metadata': self.patterns[pool_id]['metadata']
                    })

            results.sort(key=lambda x: x['similarity'], reverse=True)
            self.metrics.patterns_retrieved += 1
            return results[:k]

        except Exception as e:
            log_error(f"Pattern search failed: {e}")
            return []

    async def _cleanup_old_patterns(self):
        try:
            if time.time() - self.metrics.last_cleanup < 3600:
                return

            sorted_patterns = sorted(
                self.patterns.items(),
                key=lambda x: x[1]['timestamp']
            )

            patterns_to_remove = len(sorted_patterns) - self.config.cache_size
            for pool_id, _ in sorted_patterns[:patterns_to_remove]:
                del self.patterns[pool_id]
                del self.cache[pool_id]

            self.metrics.last_cleanup = time.time()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        except Exception as e:
            log_error(f"Pattern cleanup failed: {e}")

    async def cleanup(self):
        """Cleanup resources"""
        try:
            self.patterns.clear()
            self.cache.clear()
            self.recent_patterns.clear()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except Exception as e:
            log_error(f"Cleanup failed: {e}")

if __name__ == "__main__":
    async def test():
        async with EnhancedVectorStore() as store:
            test_vector = torch.randn(512)
            await store.store_pattern("test_pool", test_vector, {"test": True})
            results = await store.search_similar_patterns(test_vector)
            print(results)

    asyncio.run(test()) is this script able to be condensed?