import os
from typing import Literal, Sequence
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


class XQubit():
    def __init__(self, x: str):
        self.__read_exp(x)
        self.__exp = None
        
    

    def __read_exp(self, x: str) -> None:
        if not os.path.exists(x):
            raise FileNotFoundError(x)
        
        df = pd.read_csv(x, sep='\t', index_col=0)
        
        exp_mat = df.to_numpy().astype(np.float64)
        genes = df.index.to_list()
        exp_design = {
            'libraries': df.columns.to_list(),
            'samples': [],
            'replicates': []
        }
        for c in exp_design['libraries']:
            s, r = c.rsplit('_', 1)
            exp_design['samples'].append(s)
            exp_design['replicates'].append(r)
        
        # build expression cube
        uniq_s = sorted(set(exp_design['samples']))
        uniq_r = sorted(set(exp_design['replicates']))
        s2i = {s: i for i, s in enumerate(uniq_s)}
        r2i = {r: i for i, r in enumerate(uniq_r)}
        
        exp_cube = np.zeros((len(genes), len(uniq_r), len(uniq_s)), dtype=np.float64)
        for col_idx, (s, r) in enumerate(zip(exp_design['samples'], exp_design['replicates'])):
            exp_cube[:, r2i[r], s2i[s]] = exp_mat[:, col_idx]
                
        # record data
        self.__genes = genes
        self.__exp_design = exp_design
        self.__exp_design_uniq_samples = uniq_s
        self.__expcube = exp_cube
        self.__expmat = np.mean(exp_cube, axis=1)
        # initialize alpha
        self.__alphas = np.append(
            np.linspace(0, 3 * np.pi / (2 * (len(uniq_s) - 1)), 4),
            2 * np.pi / len(uniq_s))
    
  
    