database_file_manager.py

"""
Database and File Management Module for Enhanced Portfolio System
Demonstrates comprehensive database operations and file I/O
"""

import sqlite3
import json
import csv
import pandas as pd
from typing import Dict, List, Optional, Any
import os
from datetime import datetime
import pickle
import xml.etree.ElementTree as ET
from enhanced_portfolio_system import Investment, Stock, Bond, Portfolio, Investor, InvestmentType

class DatabaseManager:
    """
    Comprehensive database management class
    Handles all database operations for the portfolio system
    """
    
    def __init__(self, db_path: str = "portfolio_system.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """Initialize database with all required tables"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            
            # Create investors table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS investors (
                    investor_id TEXT PRIMARY KEY,
                    name TEXT NOT NULL,
                    email TEXT UNIQUE NOT NULL,
                    phone TEXT,
                    address TEXT,
                    registration_date TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Create portfolios table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS portfolios (
                    portfolio_id TEXT PRIMARY KEY,
                    investor_id TEXT,
                    name TEXT NOT NULL,
                    description TEXT,
                    creation_date TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (investor_id) REFERENCES investors (investor_id)
                )
            ''')
            
            # Create investments table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS investments (
                    purchase_id TEXT PRIMARY KEY,
                    portfolio_id TEXT,
                    symbol TEXT NOT NULL,
                    investment_type TEXT NOT NULL,
                    shares INTEGER NOT NULL,
                    purchase_price REAL NOT NULL,
                    current_value REAL NOT NULL,
                    purchase_date TEXT NOT NULL,
                    sector TEXT,
                    dividend_yield REAL DEFAULT 0,
                    coupon_rate REAL DEFAULT 0,
                    maturity_date TEXT,
                    credit_rating TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (portfolio_id) REFERENCES portfolios (portfolio_id)
                )
            ''')
            
            # Create market_data table for historical prices
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS market_data (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    symbol TEXT NOT NULL,
                    date TEXT NOT NULL,
                    open_price REAL,
                    high_price REAL,
                    low_price REAL,
                    close_price REAL,
                    volume INTEGER,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(symbol, date)
                )
            ''')
            
            # Create portfolio_snapshots for historical tracking
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS portfolio_snapshots (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    portfolio_id TEXT,
                    snapshot_date TEXT,
                    total_value REAL,
                    total_cost REAL,
                    total_earnings REAL,
                    investment_count INTEGER,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (portfolio_id) REFERENCES portfolios (portfolio_id)
                )
            ''')
            
            conn.commit()
    
    def save_investor(self, investor: Investor) -> bool:
        """Save investor to database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    INSERT OR REPLACE INTO investors 
                    (investor_id, name, email, phone, address, registration_date)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', (investor.investor_id, investor.name, investor.email,
                     investor.phone, investor.address, investor.registration_date))
                conn.commit()
                return True
        except Exception as e:
            print(f"Error saving investor: {e}")
            return False
    
    def save_portfolio(self, portfolio: Portfolio, investor_id: str) -> bool:
        """Save portfolio to database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    INSERT OR REPLACE INTO portfolios 
                    (portfolio_id, investor_id, name, description, creation_date)
                    VALUES (?, ?, ?, ?, ?)
                ''', (portfolio.portfolio_id, investor_id, portfolio.name,
                     portfolio.description, portfolio.creation_date))
                conn.commit()
                return True
        except Exception as e:
            print(f"Error saving portfolio: {e}")
            return False
    
    def save_investment(self, investment: Investment, portfolio_id: str) -> bool:
        """Save investment to database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # Prepare investment data
                data = [
                    investment.purchase_id, portfolio_id, investment.symbol,
                    investment.investment_type.value, investment.shares,
                    investment.purchase_price, investment.current_value,
                    investment.purchase_date
                ]
                
                # Add type-specific data
                if isinstance(investment, Stock):
                    data.extend([investment.sector, investment.dividend_yield, 0, None, None])
                elif isinstance(investment, Bond):
                    data.extend([None, 0, investment.coupon_rate, 
                               investment.maturity_date, investment.credit_rating])
                else:
                    data.extend([None, 0, 0, None, None])
                
                cursor.execute('''
                    INSERT OR REPLACE INTO investments 
                    (purchase_id, portfolio_id, symbol, investment_type, shares,
                     purchase_price, current_value, purchase_date, sector,
                     dividend_yield, coupon_rate, maturity_date, credit_rating)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', data)
                
                conn.commit()
                return True
        except Exception as e:
            print(f"Error saving investment: {e}")
            return False
    
    def load_investor(self, investor_id: str) -> Optional[Investor]:
        """Load investor from database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    SELECT investor_id, name, email, phone, address, registration_date
                    FROM investors WHERE investor_id = ?
                ''', (investor_id,))
                
                row = cursor.fetchone()
                if row:
                    investor = Investor(*row)
                    # Load portfolios
                    portfolios = self.load_portfolios_for_investor(investor_id)
                    for portfolio in portfolios:
                        investor.portfolios[portfolio.portfolio_id] = portfolio
                    return investor
                return None
        except Exception as e:
            print(f"Error loading investor: {e}")
            return None
    
    def load_portfolios_for_investor(self, investor_id: str) -> List[Portfolio]:
        """Load all portfolios for an investor"""
        portfolios = []
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    SELECT portfolio_id, name, description, creation_date
                    FROM portfolios WHERE investor_id = ?
                ''', (investor_id,))
                
                for row in cursor.fetchall():
                    portfolio = Portfolio(*row)
                    # Load investments for this portfolio
                    investments = self.load_investments_for_portfolio(row[0])
                    for investment in investments:
                        key = f"{investment.symbol}_{investment.purchase_id}"
                        portfolio.investments[key] = investment
                    portfolios.append(portfolio)
        except Exception as e:
            print(f"Error loading portfolios: {e}")
        
        return portfolios
    
    def load_investments_for_portfolio(self, portfolio_id: str) -> List[Investment]:
        """Load all investments for a portfolio"""
        investments = []
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    SELECT purchase_id, symbol, investment_type, shares,
                           purchase_price, current_value, purchase_date,
                           sector, dividend_yield, coupon_rate, maturity_date, credit_rating
                    FROM investments WHERE portfolio_id = ?
                ''', (portfolio_id,))
                
                for row in cursor.fetchall():
                    investment_type = InvestmentType(row[2])
                    
                    if investment_type == InvestmentType.STOCK:
                        investment = Stock(
                            row[0], row[1], row[3], row[4], row[5], row[6],
                            row[7] or "Unknown", row[8] or 0.0
                        )
                    elif investment_type == InvestmentType.BOND:
                        investment = Bond(
                            row[0], row[1], row[3], row[4], row[5], row[6],
                            row[9] or 0.0, row[10] or "12/31/2030", row[11] or "NR"
                        )
                    else:
                        investment = Investment(
                            row[0], row[1], row[3], row[4], row[5], row[6], investment_type
                        )
                    
                    investments.append(investment)
        except Exception as e:
            print(f"Error loading investments: {e}")
        
        return investments
    
    def save_market_data(self, market_data_list: List[Dict]) -> bool:
        """Save market data to database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                for data in market_data_list:
                    cursor.execute('''
                        INSERT OR REPLACE INTO market_data 
                        (symbol, date, open_price, high_price, low_price, close_price, volume)
                        VALUES (?, ?, ?, ?, ?, ?, ?)
                    ''', (data['symbol'], data['date'], data['open'], 
                         data['high'], data['low'], data['close'], data['volume']))
                conn.commit()
                return True
        except Exception as e:
            print(f"Error saving market data: {e}")
            return False
    
    def get_market_data(self, symbol: str, start_date: str = None, end_date: str = None) -> List[Dict]:
        """Get market data for a symbol"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                query = "SELECT * FROM market_data WHERE symbol = ?"
                params = [symbol]
                
                if start_date:
                    query += " AND date >= ?"
                    params.append(start_date)
                if end_date:
                    query += " AND date <= ?"
                    params.append(end_date)
                
                query += " ORDER BY date"
                cursor.execute(query, params)
                
                columns = [desc[0] for desc in cursor.description]
                return [dict(zip(columns, row)) for row in cursor.fetchall()]
        except Exception as e:
            print(f"Error getting market data: {e}")
            return []

class FileManager:
    """
    Comprehensive file management class
    Handles various file formats and I/O operations
    """
    
    def __init__(self, base_path: str = "/home/ubuntu/portfolio_data"):
        self.base_path = base_path
        os.makedirs(base_path, exist_ok=True)
    
    def export_portfolio_to_csv(self, portfolio: Portfolio, filename: str) -> bool:
        """Export portfolio data to CSV file"""
        try:
            filepath = os.path.join(self.base_path, filename)
            with open(filepath, 'w', newline='') as csvfile:
                fieldnames = ['symbol', 'investment_type', 'shares', 'purchase_price', 
                            'current_value', 'purchase_date', 'earnings_loss', 
                            'yearly_rate', 'total_value']
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                
                writer.writeheader()
                for investment in portfolio.investments.values():
                    row = {
                        'symbol': investment.symbol,
                        'investment_type': investment.investment_type.value,
                        'shares': investment.shares,
                        'purchase_price': investment.purchase_price,
                        'current_value': investment.current_value,
                        'purchase_date': investment.purchase_date,
                        'earnings_loss': investment.earnings_loss(),
                        'yearly_rate': investment.yearly_earnings_loss_rate(),
                        'total_value': investment.get_total_value()
                    }
                    writer.writerow(row)
            return True
        except Exception as e:
            print(f"Error exporting to CSV: {e}")
            return False
    
    def import_portfolio_from_csv(self, filename: str, portfolio: Portfolio) -> bool:
        """Import portfolio data from CSV file"""
        try:
            filepath = os.path.join(self.base_path, filename)
            with open(filepath, 'r') as csvfile:
                reader = csv.DictReader(csvfile)
                for i, row in enumerate(reader):
                    investment_type = InvestmentType(row['investment_type'])
                    purchase_id = f"IMP_{i+1:03d}"
                    
                    if investment_type == InvestmentType.STOCK:
                        investment = Stock(
                            purchase_id, row['symbol'], int(row['shares']),
                            float(row['purchase_price']), float(row['current_value']),
                            row['purchase_date']
                        )
                    elif investment_type == InvestmentType.BOND:
                        investment = Bond(
                            purchase_id, row['symbol'], int(row['shares']),
                            float(row['purchase_price']), float(row['current_value']),
                            row['purchase_date'], 4.0, "12/31/2030"
                        )
                    else:
                        investment = Investment(
                            purchase_id, row['symbol'], int(row['shares']),
                            float(row['purchase_price']), float(row['current_value']),
                            row['purchase_date'], investment_type
                        )
                    
                    portfolio.add_investment(investment)
            return True
        except Exception as e:
            print(f"Error importing from CSV: {e}")
            return False
    
    def export_portfolio_to_json(self, portfolio: Portfolio, filename: str) -> bool:
        """Export portfolio data to JSON file"""
        try:
            filepath = os.path.join(self.base_path, filename)
            portfolio_data = portfolio.to_dict()
            with open(filepath, 'w') as jsonfile:
                json.dump(portfolio_data, jsonfile, indent=2, default=str)
            return True
        except Exception as e:
            print(f"Error exporting to JSON: {e}")
            return False
    
    def export_investor_to_json(self, investor: Investor, filename: str) -> bool:
        """Export complete investor data to JSON file"""
        try:
            filepath = os.path.join(self.base_path, filename)
            investor_data = {
                'investor_info': investor.get_investor_summary()['investor_info'],
                'portfolios': {pid: portfolio.to_dict() 
                             for pid, portfolio in investor.portfolios.items()}
            }
            with open(filepath, 'w') as jsonfile:
                json.dump(investor_data, jsonfile, indent=2, default=str)
            return True
        except Exception as e:
            print(f"Error exporting investor to JSON: {e}")
            return False
    
    def save_portfolio_snapshot(self, portfolio: Portfolio, filename: str) -> bool:
        """Save portfolio snapshot using pickle for complete object serialization"""
        try:
            filepath = os.path.join(self.base_path, filename)
            with open(filepath, 'wb') as picklefile:
                pickle.dump(portfolio, picklefile)
            return True
        except Exception as e:
            print(f"Error saving portfolio snapshot: {e}")
            return False
    
    def load_portfolio_snapshot(self, filename: str) -> Optional[Portfolio]:
        """Load portfolio snapshot from pickle file"""
        try:
            filepath = os.path.join(self.base_path, filename)
            with open(filepath, 'rb') as picklefile:
                return pickle.load(picklefile)
        except Exception as e:
            print(f"Error loading portfolio snapshot: {e}")
            return None
    
    def generate_portfolio_report(self, portfolio: Portfolio, filename: str) -> bool:
        """Generate comprehensive portfolio report"""
        try:
            filepath = os.path.join(self.base_path, filename)
            summary = portfolio.get_portfolio_summary()
            
            with open(filepath, 'w') as reportfile:
                reportfile.write("PORTFOLIO PERFORMANCE REPORT\n")
                reportfile.write("=" * 50 + "\n\n")
                
                # Portfolio info
                info = summary['portfolio_info']
                reportfile.write(f"Portfolio: {info['name']}\n")
                reportfile.write(f"ID: {info['id']}\n")
                reportfile.write(f"Description: {info['description']}\n")
                reportfile.write(f"Created: {info['creation_date']}\n\n")
                
                # Basic metrics
                metrics = summary['basic_metrics']
                reportfile.write("FINANCIAL SUMMARY\n")
                reportfile.write("-" * 20 + "\n")
                reportfile.write(f"Total Value: ${metrics['total_value']:,.2f}\n")
                reportfile.write(f"Total Cost: ${metrics['total_cost']:,.2f}\n")
                reportfile.write(f"Total Earnings: ${metrics['total_earnings']:,.2f}\n")
                reportfile.write(f"Return Percentage: {metrics['total_return_percentage']:.2f}%\n")
                reportfile.write(f"Average Yearly Return: {metrics['average_yearly_return']:.2f}%\n\n")
                
                # Holdings details
                reportfile.write("INDIVIDUAL HOLDINGS\n")
                reportfile.write("-" * 20 + "\n")
                reportfile.write(f"{'Symbol':<10}{'Type':<10}{'Shares':<8}{'Value':<12}{'Earnings':<12}{'Yearly %':<10}\n")
                reportfile.write("-" * 70 + "\n")
                
                for investment in portfolio.investments.values():
                    reportfile.write(f"{investment.symbol:<10}")
                    reportfile.write(f"{investment.investment_type.value:<10}")
                    reportfile.write(f"{investment.shares:<8}")
                    reportfile.write(f"${investment.get_total_value():<11,.2f}")
                    reportfile.write(f"${investment.earnings_loss():<11,.2f}")
                    reportfile.write(f"{investment.yearly_earnings_loss_rate():<9.2f}%\n")
                
                # Risk metrics
                if 'risk_metrics' in summary:
                    risk = summary['risk_metrics']
                    reportfile.write(f"\nRISK ANALYSIS\n")
                    reportfile.write("-" * 15 + "\n")
                    reportfile.write(f"Volatility: {risk['volatility']:.2f}%\n")
                    reportfile.write(f"Sharpe Ratio: {risk['sharpe_ratio']:.2f}\n")
                    reportfile.write(f"Max Return: {risk['max_return']:.2f}%\n")
                    reportfile.write(f"Min Return: {risk['min_return']:.2f}%\n")
            
            return True
        except Exception as e:
            print(f"Error generating report: {e}")
            return False

# Utility functions for data processing
def process_market_data_file(filepath: str) -> List[Dict]:
    """Process market data from various file formats"""
    data = []
    try:
        if filepath.endswith('.csv'):
            df = pd.read_csv(filepath)
            data = df.to_dict('records')
        elif filepath.endswith('.json'):
            with open(filepath, 'r') as f:
                data = json.load(f)
        elif filepath.endswith('.xlsx'):
            df = pd.read_excel(filepath)
            data = df.to_dict('records')
    except Exception as e:
        print(f"Error processing market data file: {e}")
    
    return data

def batch_update_prices(db_manager: DatabaseManager, price_updates: Dict[str, float]) -> int:
    """Batch update current prices for multiple symbols"""
    updated_count = 0
    try:
        with sqlite3.connect(db_manager.db_path) as conn:
            cursor = conn.cursor()
            for symbol, new_price in price_updates.items():
                cursor.execute('''
                    UPDATE investments 
                    SET current_value = ?, updated_at = CURRENT_TIMESTAMP
                    WHERE symbol = ?
                ''', (new_price, symbol))
                updated_count += cursor.rowcount
            conn.commit()
    except Exception as e:
        print(f"Error in batch price update: {e}")
    
    return updated_count

if __name__ == "__main__":
    print("Database and File Management Module - Testing")
    print("=" * 50)
    
    # Test database operations
    db_manager = DatabaseManager("test_portfolio.db")
    file_manager = FileManager()
    
    # Create test data
    from enhanced_portfolio_system import Investor, Portfolio, Stock, Bond
    
    investor = Investor("TEST001", "Test User", "test@email.com")
    portfolio = investor.create_portfolio("TEST_PORT", "Test Portfolio")
    
    stock = Stock("STK001", "AAPL", 100, 150.0, 175.0, "01/15/2024", "Technology", 0.5)
    bond = Bond("BND001", "US10Y", 10, 1000.0, 980.0, "03/10/2024", 4.5, "03/10/2034", "AAA")
    
    portfolio.add_investment(stock)
    portfolio.add_investment(bond)
    
    # Test database operations
    print("Testing database operations...")
    db_manager.save_investor(investor)
    db_manager.save_portfolio(portfolio, investor.investor_id)
    
    for investment in portfolio.investments.values():
        db_manager.save_investment(investment, portfolio.portfolio_id)
    
    # Test file operations
    print("Testing file operations...")
    file_manager.export_portfolio_to_csv(portfolio, "test_portfolio.csv")
    file_manager.export_portfolio_to_json(portfolio, "test_portfolio.json")
    file_manager.generate_portfolio_report(portfolio, "test_report.txt")
    
    print("Database and file operations completed successfully!")