import pandas as pd
import sqlite3
from pathlib import Path
import traceback
import math
import numpy as np
import re
import random
from datetime import datetime
from collections import Counter
import sys
import os

# Pārbauda un importē nepieciešamās bibliotēkas
try:
    from tqdm import tqdm
    tqdm.pandas(desc="Aprēķina datus")
    
    from pandarallel import pandarallel
    from multiprocessing import Pool, Manager, cpu_count
except ImportError:
    print("KĻŪDA: Trūkst nepieciešamās bibliotēkas. Lūdzu, instalējiet tās:")
    print("pip3 install tqdm pandarallel")
    sys.exit(1)

# --- 1. KONFIGURĀCIJA UN CEĻI ---
SCRIPT_DIR = Path(__file__).parent.resolve()
PROJECT_ROOT = SCRIPT_DIR.parent

# Mape ar CSV datiem
CSV_DIR = SCRIPT_DIR / "csv"

# Izejas mapes konfigurācija
OUTPUT_DIR = SCRIPT_DIR / "out" / "nace"
OUTPUT_DB_FILE = OUTPUT_DIR / "katalogs.sqlite"

NACE_FILE = SCRIPT_DIR / 'NACE.csv'

# CSV failu definīcijas
REGISTER_FILE = CSV_DIR / 'register.csv'
TAX_DATA_FILE = CSV_DIR / 'pdb_nm_komersantu_samaksato_nodoklu_kopsumas_odata.csv'
FINANCIAL_STATEMENTS_FILE = CSV_DIR / 'financial_statements.csv'
INCOME_STATEMENTS_FILE = CSV_DIR / 'income_statements.csv'
BALANCE_SHEETS_FILE = CSV_DIR / 'balance_sheets.csv'

# Mēģina importēt algas kalkulatoru
PYTHON_GENERATOR_DIR = PROJECT_ROOT / 'python_generator'
if PYTHON_GENERATOR_DIR.is_dir():
    sys.path.insert(0, str(PROJECT_ROOT))
    sys.path.insert(0, str(PYTHON_GENERATOR_DIR))

try:
    from core.page_builder import _calculate_net_salary
except ImportError:
    def _calculate_net_salary(gross, year):
        return {'net_salary': gross * 0.70}

# --- 2. PALĪGFUNKCIJAS ---

def clean_id(val):
    """Palīgfunkcija ID tīrīšanai: noņem .0 un atstarpes"""
    if pd.isna(val): return ""
    s = str(val).strip()
    if s.endswith('.0'):
        return s[:-2]
    return s

def calculate_financial_ratios(income_data, balance_data):
    ratios = {}
    i = {}
    b = {}
    def get_val(series, key):
        if series is None or key not in series: return 0
        val = pd.to_numeric(series.get(key), errors='coerce')
        return 0 if pd.isna(val) else val
    
    if income_data is not None: i = {k: get_val(income_data, k) for k in income_data.keys()}
    if balance_data is not None: b = {k: get_val(balance_data, k) for k in balance_data.keys()}
    
    def safe_divide(numerator, denominator):
        if denominator == 0: return None
        try:
            result = numerator / denominator
            return result if math.isfinite(result) else None
        except TypeError: return None

    ratios['current_ratio'] = safe_divide(b.get('total_current_assets', 0), b.get('current_liabilities', 0))
    ratios['quick_ratio'] = safe_divide(b.get('total_current_assets', 0) - b.get('inventories', 0), b.get('current_liabilities', 0))
    ratios['net_profit_margin'] = safe_divide(i.get('net_income', 0), i.get('net_turnover', 0))
    if ratios['net_profit_margin'] is not None: ratios['net_profit_margin'] *= 100
    ratios['roa'] = safe_divide(i.get('net_income', 0), b.get('total_assets', 0))
    if ratios['roa'] is not None: ratios['roa'] *= 100
    ratios['roe'] = safe_divide(i.get('net_income', 0), b.get('equity', 0))
    if ratios['roe'] is not None: ratios['roe'] *= 100
    
    total_liabilities = b.get('current_liabilities', 0) + b.get('non_current_liabilities', 0)
    ratios['debt_to_equity'] = safe_divide(total_liabilities, b.get('equity', 0))
    ratios['asset_turnover'] = safe_divide(i.get('net_turnover', 0), b.get('total_assets', 0))
    
    ebit = i.get('income_before_income_taxes', 0) + i.get('interest_expenses', 0)
    iegulditais_kapitals = b.get('total_assets', 0) - b.get('current_liabilities', 0)
    ratios['roce'] = safe_divide(ebit, iegulditais_kapitals)
    if ratios['roce'] is not None: ratios['roce'] *= 100
    ratios['interest_coverage'] = safe_divide(ebit, i.get('interest_expenses', 0))
    
    try:
        A = safe_divide(b.get('total_current_assets', 0) - b.get('current_liabilities', 0), b.get('total_assets', 0))
        B = safe_divide(b.get('equity', 0), b.get('total_assets', 0))
        C = safe_divide(ebit, b.get('total_assets', 0))
        D = safe_divide(b.get('equity', 0), total_liabilities)
        E = safe_divide(i.get('net_turnover', 0), b.get('total_assets', 0))
        
        if None in [A, B, C, D, E] or not all(math.isfinite(v) for v in [A, B, C, D, E] if v is not None):
            ratios['altman_z_score'] = None
        else:
            # Altman Z' (1983) modelis privātiem uzņēmumiem
            ratios['altman_z_score'] = 0.717 * A + 0.847 * B + 3.107 * C + 0.420 * D + 0.998 * E
    except (TypeError, ZeroDivisionError):
        ratios['altman_z_score'] = None
    return ratios

def _get_health_symbol(ratio_key, value):
    if value is None or not math.isfinite(value): return 'n'
    if ratio_key == 'current_ratio':
        return 'r' if value < 1.0 else 'o' if value < 1.5 else 'g' if value <= 3.0 else 'b'
    if ratio_key == 'quick_ratio':
        return 'r' if value < 0.8 else 'o' if value < 1.0 else 'g' if value <= 2.0 else 'b'
    if ratio_key == 'net_profit_margin':
        return 'r' if value < 0 else 'o' if value < 10 else 'g'
    if ratio_key == 'roa':
        return 'r' if value < 2 else 'o' if value <= 5 else 'g'
    if ratio_key == 'roe':
        return 'r' if value < 5 else 'o' if value <= 15 else 'g'
    if ratio_key == 'debt_to_equity':
        return 'g' if value < 1.5 else 'o' if value <= 2.0 else 'r'
    if ratio_key == 'asset_turnover':
        return 'r' if value < 0.5 else 'o' if value <= 1.0 else 'g'
    if ratio_key == 'roce':
        return 'r' if value < 5 else 'o' if value <= 15 else 'g'
    if ratio_key == 'interest_coverage':
        return 'r' if value < 1.5 else 'o' if value < 3.0 else 'g'
    if ratio_key == 'altman_z_score':
        # 1983. gada modeļa sliekšņi
        return 'r' if value < 1.23 else 'o' if value < 2.90 else 'g'
    return 'n'

def calculate_pct_change(current, previous):
    """Aprēķina % izmaiņu starp diviem skaitļiem."""
    if pd.isna(current) or pd.isna(previous) or previous == 0:
        return None
    try:
        change = ((current - previous) / abs(previous)) * 100
        return change if math.isfinite(change) else None
    except:
        return None

# --- 3. DATUBĀZES INITIALIZĀCIJA ---

def create_database():
    if not OUTPUT_DIR.exists():
        print(f"Izveido mapi: {OUTPUT_DIR}")
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    print(f"Izveido/pārraksta datubāzi: {OUTPUT_DB_FILE}")
    if OUTPUT_DB_FILE.exists():
        try:
            OUTPUT_DB_FILE.unlink()
        except PermissionError:
            print("KĻŪDA: Nevar izdzēst veco DB failu. Iespējams, tas ir atvērts citā programmā.")
            sys.exit(1)
            
    conn = sqlite3.connect(OUTPUT_DB_FILE)
    cursor = conn.cursor()
    try:
        cursor.execute('''
            CREATE TABLE nace (
                code TEXT PRIMARY KEY,
                name TEXT,
                parent_code TEXT,
                level INTEGER,
                health_string TEXT
            )
        ''')
        
        # Pievienotas _change un prev_ kolonnas
        cursor.execute('''
            CREATE TABLE companies (
                regcode TEXT PRIMARY KEY,
                name TEXT,
                nace_code_np TEXT,
                turnover REAL,
                turnover_change REAL,
                prev_turnover REAL,
                employees INTEGER,
                employees_change REAL,
                prev_employees INTEGER,
                profit REAL,
                profit_change REAL,
                prev_profit REAL,
                avg_gross_salary REAL,
                avg_net_salary REAL,
                salary_change REAL,
                prev_salary REAL,
                registered_date TEXT,
                financial_health_string TEXT,
                location TEXT 
            )
        ''')
        
        cursor.execute('CREATE INDEX idx_nace_parent ON nace (parent_code)')
        cursor.execute('CREATE INDEX idx_companies_nace ON companies (nace_code_np)')
        conn.commit()
        print("  Tabulas un indeksi izveidoti.")
    except Exception as e:
        print(f"Kļūda, veidojot datubāzes struktūru: {e}")
        conn.rollback()
    finally:
        conn.close()

# --- 4. DATU IELĀDE UN APSTRĀDE ---

def load_nace_data():
    print("Ielādē datus no NACE.csv...")
    if not NACE_FILE.exists():
        print(f"KĻŪDA: 'NACE.csv' nav atrasts: {NACE_FILE}")
        return False

    try:
        nace_df = pd.read_csv(
            NACE_FILE, 
            dtype={'Kods': str, 'Vecāka kods': str, 'Līmenis': str}, 
            usecols=['Kods', 'Nosaukums', 'Vecāka kods', 'Līmenis'],
            encoding='utf-8-sig'
        )
        
        nace_df['Kods'] = nace_df['Kods'].str.strip()
        nace_df['Nosaukums'] = nace_df['Nosaukums'].str.strip()
        nace_df['Vecāka kods'] = nace_df['Vecāka kods'].fillna('').str.strip()
        nace_df['Līmenis'] = pd.to_numeric(nace_df['Līmenis'], errors='coerce').fillna(9).astype(int)
        nace_df['health_string'] = 'nnnnnnnnnn'

        conn = sqlite3.connect(OUTPUT_DB_FILE)
        cursor = conn.cursor()
        
        data_to_insert = list(nace_df[['Kods', 'Nosaukums', 'Vecāka kods', 'Līmenis', 'health_string']].itertuples(index=False, name=None))
        
        cursor.executemany(
            "INSERT OR IGNORE INTO nace (code, name, parent_code, level, health_string) VALUES (?, ?, ?, ?, ?)",
            data_to_insert
        )
        conn.commit()
        print(f"  Ievietoti {len(data_to_insert)} NACE ieraksti.")
        conn.close()
        return True
    except Exception as e:
        print(f"Kļūda, apstrādājot NACE.csv: {e}")
        return False

def load_and_prepare_company_data():
    print("\nSāk uzņēmumu datu apstrādi...")
    
    required_files = [REGISTER_FILE, TAX_DATA_FILE, FINANCIAL_STATEMENTS_FILE, INCOME_STATEMENTS_FILE, BALANCE_SHEETS_FILE]
    for f in required_files:
        if not f.exists():
            print(f"KĻŪDA: Trūkst fails: {f}")
            return pd.DataFrame()

    # 1. Reģistrs
    print("  Ielasa Reģistru...")
    register_df = pd.read_csv(
        REGISTER_FILE, delimiter=';',
        usecols=['regcode', 'name', 'closed', 'terminated', 'type', 'name_in_quotes', 'registered', 'address'],
        dtype=str,
        encoding='utf-8-sig'
    )
    
    register_df['closed'] = register_df['closed'].fillna('')
    register_df['terminated'] = register_df['terminated'].fillna('')
    active_mask = (register_df['closed'] != 'L') & (register_df['closed'] != 'R') & \
                  ((register_df['terminated'] == '') | (register_df['terminated'] == '0000-00-00'))
    active_register_df = register_df[active_mask].copy()
    active_register_df['regcode'] = active_register_df['regcode'].str.strip()
    active_register_df['location'] = active_register_df['address'].fillna('').str.split(',', n=1, expand=False).str[0].str.strip()
    
    tqdm.pandas(desc="  Formatē nosaukumus")
    def format_name(row):
        t = str(row.get('type', '')).strip()
        n = str(row.get('name_in_quotes', '')).strip()
        if t and n: return f"{t} \"{n}\""
        elif n: return f"\"{n}\""
        return str(row.get('name', '')).strip()
    
    active_register_df['name'] = active_register_df.progress_apply(format_name, axis=1)
    active_register_df = active_register_df[['regcode', 'name', 'registered', 'location']]

    # 2. Nodokļi (Darbinieki un Algas)
    print("  Ielasa un apstrādā Nodokļu datus (izmaiņu aprēķins)...")
    tax_cols = ['Registracijas_kods', 'Pamatdarbibas_NACE_kods', 'Taksacijas_gads', 'Taja_skaita_VSAOI', 'Videjais_nodarbinato_personu_skaits_cilv']
    try:
        tax_df = pd.read_csv(TAX_DATA_FILE, delimiter=';', usecols=tax_cols, dtype=str, on_bad_lines='skip', encoding='utf-8-sig')
    except ValueError:
        tax_df = pd.read_csv(TAX_DATA_FILE, delimiter=',', usecols=tax_cols, dtype=str, on_bad_lines='skip', encoding='utf-8-sig')

    tax_df['regcode'] = tax_df['Registracijas_kods'].str.strip()
    tax_df['Taksacijas_gads'] = pd.to_numeric(tax_df['Taksacijas_gads'], errors='coerce').fillna(0).astype(int)
    tax_df['Videjais_nodarbinato_personu_skaits_cilv'] = pd.to_numeric(tax_df['Videjais_nodarbinato_personu_skaits_cilv'], errors='coerce').fillna(0)
    
    # Tīram VSAOI un aprēķinam Bruto algu Visiem gadiem (lai varētu salīdzināt)
    tax_df['vsaoi_clean'] = tax_df['Taja_skaita_VSAOI'].astype(str).str.replace(r'\s+', '', regex=True).str.replace(',', '.')
    tax_df['vsaoi_val'] = pd.to_numeric(tax_df['vsaoi_clean'], errors='coerce').fillna(0) * 1000
    
    # Bruto algas aprēķins (vektorizēts)
    tax_df['gross_salary_calc'] = np.where(
        (tax_df['Videjais_nodarbinato_personu_skaits_cilv'] > 0) & (tax_df['vsaoi_val'] > 0),
        (tax_df['vsaoi_val'] / 0.3409) / tax_df['Videjais_nodarbinato_personu_skaits_cilv'] / 12,
        0
    )
    
    # Neto algas aprēķins (vektorizēts, aptuvens priekš salīdzināšanas)
    # Izmantojam vienkāršotu 70% pieņēmumu, jo _calculate_net_salary ir lēns un Python funkcija
    tax_df['net_salary_calc'] = tax_df['gross_salary_calc'] * 0.70

    # NACE kods
    tax_df['Pamatdarbibas_NACE_kods'] = tax_df['Pamatdarbibas_NACE_kods'].str.strip().replace('?', np.nan)
    tax_df.sort_values(by=['regcode', 'Taksacijas_gads'], ascending=[True, True], inplace=True)
    tax_df['Pamatdarbibas_NACE_kods'] = tax_df.groupby('regcode')['Pamatdarbibas_NACE_kods'].ffill()
    tax_df['Pamatdarbibas_NACE_kods'] = tax_df['Pamatdarbibas_NACE_kods'].str.replace('.', '', regex=False).str.strip()

    # --- Aprēķinam izmaiņas un saglabājam prev vērtības ---
    grouped_tax = tax_df.groupby('regcode')
    tax_df['prev_employees'] = grouped_tax['Videjais_nodarbinato_personu_skaits_cilv'].shift(1)
    tax_df['prev_net_salary'] = grouped_tax['net_salary_calc'].shift(1) # Neto algas prev vērtība
    
    # Darbinieku izmaiņas %
    tax_df['employees_change'] = np.where(
        (tax_df['prev_employees'] > 0),
        ((tax_df['Videjais_nodarbinato_personu_skaits_cilv'] - tax_df['prev_employees']) / tax_df['prev_employees']) * 100,
        None
    )
    # Algu izmaiņas % (Neto algas)
    tax_df['salary_change'] = np.where(
        (tax_df['prev_net_salary'] > 0),
        ((tax_df['net_salary_calc'] - tax_df['prev_net_salary']) / tax_df['prev_net_salary']) * 100,
        None
    )

    # Paturam tikai jaunāko gadu katram uzņēmumam
    tax_df_latest = tax_df.sort_values(by='Taksacijas_gads', ascending=False).drop_duplicates(subset=['regcode'])
    
    company_info_map = tax_df_latest[['regcode', 'Pamatdarbibas_NACE_kods', 'Taksacijas_gads', 'Taja_skaita_VSAOI', 
                                      'Videjais_nodarbinato_personu_skaits_cilv', 'employees_change', 'prev_employees', 
                                      'salary_change', 'prev_net_salary']]
    company_info_map.rename(columns={'prev_net_salary': 'prev_salary'}, inplace=True)

    # 3. Finanšu dati (Apgrozījums un Peļņa)
    print("  Ielasa un apstrādā Finanšu pārskatus (izmaiņu aprēķins)...")
    fs_df = pd.read_csv(
        FINANCIAL_STATEMENTS_FILE, delimiter=';',
        usecols=['legal_entity_registration_number', 'year', 'id', 'employees', 'rounded_to_nearest', 'source_type'],
        dtype=str,
        encoding='utf-8-sig'
    )
    fs_df.rename(columns={'id': 'statement_id', 'legal_entity_registration_number': 'regcode'}, inplace=True)
    fs_df['regcode'] = fs_df['regcode'].str.strip()
    fs_df['statement_id'] = fs_df['statement_id'].apply(clean_id)
    fs_df['year'] = pd.to_numeric(fs_df['year'], errors='coerce')

    # Ielasa finanšu datus (Turnover, Net Income, Assets, etc.)
    income_df = pd.read_csv(INCOME_STATEMENTS_FILE, delimiter=';', dtype=str, encoding='utf-8-sig')
    income_df['statement_id'] = income_df['statement_id'].apply(clean_id)
    income_df = income_df.set_index('statement_id')
    
    bal_cols = pd.read_csv(BALANCE_SHEETS_FILE, delimiter=';', nrows=0, encoding='utf-8-sig').columns
    use_cols_bal = [c for c in bal_cols if c.lower() != 'file_id']
    balance_df = pd.read_csv(BALANCE_SHEETS_FILE, delimiter=';', dtype=str, usecols=use_cols_bal, encoding='utf-8-sig')
    balance_df['statement_id'] = balance_df['statement_id'].apply(clean_id)
    balance_df = balance_df.set_index('statement_id')

    # Savieno visu kopā
    fs_full = fs_df.join(income_df, on='statement_id').join(balance_df, on='statement_id')
    
    # Konvertē skaitļus
    def get_factor(rounding_str):
        r = str(rounding_str).upper().strip()
        return 1000 if r == 'THOUSANDS' else 1000000 if r == 'MILLIONS' else 1
    
    fs_full['factor'] = fs_full['rounded_to_nearest'].apply(get_factor)
    
    cols_to_convert = ['net_income', 'net_turnover', 'total_current_assets', 'inventories', 
                       'current_liabilities', 'non_current_liabilities', 'equity', 
                       'total_assets', 'interest_expenses', 'income_before_income_taxes']
    
    for col in cols_to_convert:
        if col in fs_full.columns:
            fs_full[col] = pd.to_numeric(fs_full[col], errors='coerce') * fs_full['factor']

    # --- Izmaiņu aprēķins ---
    # Prioritizējam UGP virs UKGP dublikātu gadījumā (viena gada ietvaros)
    fs_full['priority'] = fs_full['source_type'].apply(lambda x: 1 if str(x).upper().strip() == 'UGP' else 2)
    fs_full.sort_values(by=['regcode', 'year', 'priority'], ascending=[True, True, True], inplace=True)
    
    # Atstājam tikai vienu ierakstu gadā katram uzņēmumam (labāko)
    fs_yearly = fs_full.drop_duplicates(subset=['regcode', 'year'], keep='first').copy()
    
    grouped_fs = fs_yearly.groupby('regcode')
    fs_yearly['prev_turnover'] = grouped_fs['net_turnover'].shift(1)
    fs_yearly['prev_profit'] = grouped_fs['net_income'].shift(1)
    
    fs_yearly['turnover_change'] = np.where(
        (fs_yearly['prev_turnover'].abs() > 0), 
        ((fs_yearly['net_turnover'] - fs_yearly['prev_turnover']) / fs_yearly['prev_turnover'].abs()) * 100,
        None
    )
    
    fs_yearly['profit_change'] = np.where(
        (fs_yearly['prev_profit'].abs() > 0),
        ((fs_yearly['net_income'] - fs_yearly['prev_profit']) / fs_yearly['prev_profit'].abs()) * 100,
        None
    )

    # Ņemam tikai jaunāko gadu
    latest_fs_df = fs_yearly.sort_values(by='year', ascending=False).drop_duplicates(subset=['regcode'])
    
    latest_fs_df['profit'] = latest_fs_df.get('net_income', np.nan)
    latest_fs_df['turnover'] = latest_fs_df.get('net_turnover', np.nan)
    latest_fs_df['employees_fs'] = pd.to_numeric(latest_fs_df['employees'], errors='coerce').fillna(0).astype(int)

    # Apvieno tabulas
    print("  Apvieno datus un saglabā DB...")
    merged_df = pd.merge(active_register_df, company_info_map, on='regcode', how='left')
    merged_df = pd.merge(merged_df, latest_fs_df, on='regcode', how='left')
    
    # Tīrīšana
    merged_df = merged_df.loc[:, ~merged_df.columns.duplicated()]

    # Darbinieku skaits (prioritāte FS)
    merged_df['employees'] = np.where(
        merged_df['employees_fs'] > 0, 
        merged_df['employees_fs'], 
        merged_df['Videjais_nodarbinato_personu_skaits_cilv']
    )
    merged_df['employees'] = pd.to_numeric(merged_df['employees'], errors='coerce').fillna(0).astype(int)

    # --- PARALĒLĀ APSTRĀDE (Algas un Rādītāji) ---
    print("  Sākas paralēlā datu apstrāde...")
    pandarallel.initialize(nb_workers=cpu_count(), progress_bar=True, verbose=0)

    def calculate_extra_data(row):
        avg_gross_salary, avg_net_salary = None, None
        financial_health_string = 'nnnnnnnnnn'
        try:
            # Algas aprēķins te paliek tikai priekš 'avg_net_salary' precīzas vērtības
            # Izmaiņas % jau ir aprēķināts iepriekš
            vsaoi_str = re.sub(r'\s+', '', str(row.get('Taja_skaita_VSAOI', ''))).replace(',', '.')
            vsaoi = pd.to_numeric(vsaoi_str, errors='coerce') * 1000
            employees_tax = pd.to_numeric(row.get('Videjais_nodarbinato_personu_skaits_cilv'), errors='coerce')

            if employees_tax > 0 and vsaoi > 0:
                avg_gross_salary = (vsaoi / 0.3409) / employees_tax / 12
                year_val = row.get('Taksacijas_gads')
                year = int(year_val) if pd.notna(year_val) and year_val != 0 else datetime.now().year
                avg_net_salary = _calculate_net_salary(avg_gross_salary, year).get('net_salary')
        except: pass

        try:
            ratios = calculate_financial_ratios(row, row)
            ks = ['current_ratio', 'quick_ratio', 'net_profit_margin', 'roa', 'roe',
                  'debt_to_equity', 'asset_turnover', 'roce', 'interest_coverage', 'altman_z_score']
            financial_health_string = "".join([_get_health_symbol(k, ratios.get(k)) for k in ks])
        except: pass
        
        return pd.Series([avg_gross_salary, avg_net_salary, financial_health_string])

    results_df = merged_df.parallel_apply(calculate_extra_data, axis=1)
    results_df.columns = ['avg_gross_salary', 'avg_net_salary', 'financial_health_string']
    
    final_df = pd.concat([merged_df, results_df], axis=1)
    
    final_df['nace_code_np'] = final_df['Pamatdarbibas_NACE_kods'].fillna('UNDEFINED')

    # Kolonnu karte uz DB
    db_cols = {
        'regcode': 'regcode', 
        'name': 'name', 
        'nace_code_np': 'nace_code_np', 
        'turnover': 'turnover',
        'turnover_change': 'turnover_change',
        'prev_turnover': 'prev_turnover',
        'employees': 'employees',
        'employees_change': 'employees_change',
        'prev_employees': 'prev_employees',
        'profit': 'profit',
        'profit_change': 'profit_change',
        'prev_profit': 'prev_profit',
        'avg_gross_salary': 'avg_gross_salary', 
        'avg_net_salary': 'avg_net_salary',
        'salary_change': 'salary_change',
        'prev_salary': 'prev_salary',
        'registered': 'registered_date', 
        'financial_health_string': 'financial_health_string',
        'location': 'location'
    }
    
    # Sagatavo saglabāšanai
    source_keys = list(db_cols.keys())
    for key in source_keys:
        if key not in final_df.columns:
            final_df[key] = None
            
    final_df_clean = final_df[source_keys].copy()
    final_df_clean.rename(columns=db_cols, inplace=True)
    final_df_clean = final_df_clean.where(pd.notnull(final_df_clean), None)
    
    final_df_clean['employees'] = pd.to_numeric(final_df_clean['employees']).fillna(0).astype(int)

    print(f"  Sagatavoti {len(final_df_clean)} ieraksti.")
    return final_df_clean

def insert_companies(df):
    if df is None or df.empty: return
    print(f"Ievieto {len(df)} uzņēmumus datubāzē...")
    conn = sqlite3.connect(OUTPUT_DB_FILE)
    records = df.to_dict('records')
    sql = """INSERT INTO companies (
        regcode, name, nace_code_np, 
        turnover, turnover_change, prev_turnover,
        employees, employees_change, prev_employees,
        profit, profit_change, prev_profit,
        avg_gross_salary, avg_net_salary, salary_change, prev_salary,
        registered_date, financial_health_string, location
    ) VALUES (
        :regcode, :name, :nace_code_np, 
        :turnover, :turnover_change, :prev_turnover,
        :employees, :employees_change, :prev_employees,
        :profit, :profit_change, :prev_profit,
        :avg_gross_salary, :avg_net_salary, :salary_change, :prev_salary,
        :registered_date, :financial_health_string, :location
    )"""
    cursor = conn.cursor()
    try:
        cursor.executemany(sql, records)
        conn.commit()
        print("  Dati veiksmīgi saglabāti.")
    except Exception as e: print(f"Kļūda saglabājot datus: {e}")
    finally: conn.close()

# --- 5. NACE VESELĪBAS APRĒĶINS ---
GLOBAL_COMPANIES_DF = None
GLOBAL_NACE_MAP = None

def init_worker(df_list, cols, shared_map):
    global GLOBAL_COMPANIES_DF, GLOBAL_NACE_MAP
    GLOBAL_COMPANIES_DF = pd.DataFrame(df_list, columns=cols)
    GLOBAL_NACE_MAP = shared_map

def process_nace_health(nace_code):
    if nace_code in GLOBAL_NACE_MAP: return GLOBAL_NACE_MAP[nace_code]
    stripped_code = nace_code.replace('.', '')
    if 'nace_code_np' not in GLOBAL_COMPANIES_DF.columns: return 'nnnnnnnnnn'
    subset = GLOBAL_COMPANIES_DF[GLOBAL_COMPANIES_DF['nace_code_np'].str.startswith(stripped_code, na=False)]
    if subset.empty: res = 'nnnnnnnnnn'
    else:
        res = ""
        for i in range(10):
            col_chars = [s[i] for s in subset['financial_health_string'] if len(s) == 10 and s[i] != 'n']
            if not col_chars: res += 'n'
            else: res += Counter(col_chars).most_common(1)[0][0]
    GLOBAL_NACE_MAP[nace_code] = res
    return res

def calculate_nace_health_strings():
    print("\nAprēķina NACE nozaru veselības rādītājus...")
    conn = sqlite3.connect(OUTPUT_DB_FILE)
    nace_codes = [r[0] for r in conn.execute("SELECT code FROM nace").fetchall()]
    if 'UNDEFINED' not in nace_codes:
        nace_codes.append('UNDEFINED')
        conn.execute("INSERT OR IGNORE INTO nace (code, name, level, health_string) VALUES ('UNDEFINED', 'Nedefinēta nozare', 1, 'nnnnnnnnnn')")
        conn.commit()
    df = pd.read_sql("SELECT nace_code_np, financial_health_string FROM companies WHERE financial_health_string != 'nnnnnnnnnn'", conn)
    conn.close()
    if df.empty:
        print("  Nav datu priekš aprēķina.")
        return
    with Manager() as manager:
        shared_map = manager.dict()
        df_list = df.values.tolist()
        df_cols = df.columns.tolist()
        with Pool(cpu_count(), initializer=init_worker, initargs=(df_list, df_cols, shared_map)) as pool:
            results = list(tqdm(pool.imap(process_nace_health, nace_codes), total=len(nace_codes), desc="  Agregē datus"))
        conn = sqlite3.connect(OUTPUT_DB_FILE)
        updates = list(zip(results, nace_codes))
        conn.executemany("UPDATE nace SET health_string = ? WHERE code = ?", updates)
        conn.commit()
        conn.close()
        print(f"  Atjaunotas {len(updates)} nozares.")

# --- 6. GALVENĀ IZPILDE ---
if __name__ == "__main__":
    os.chdir(SCRIPT_DIR)
    start_time = datetime.now()
    print("=== SĀK DATUBĀZES ĢENERĒŠANU ===")
    print(f"Datu mape: {CSV_DIR}")
    create_database()
    if load_nace_data():
        companies_df = load_and_prepare_company_data()
        if not companies_df.empty:
            insert_companies(companies_df)
            calculate_nace_health_strings()
    print(f"\nDARBS PABEIGTS {datetime.now() - start_time}")
    print(f"Rezultāts saglabāts: {OUTPUT_DB_FILE}")
