use numpy::{PyArray, PyArray1, PyArray2};
use pyo3::exceptions::PyTypeError;
use pyo3::prelude::*;
use pyo3::types::PySlice;

#[pyclass(sequence)]
pub struct DnaSequence {
    pub length: usize,  // Store the number of nucleotides
    pub data: Vec<u32>, // Store the encoded data
}

const fn create_nucleotide_lookup() -> [[u8; 2]; 256] {
    let mut table = [[b'N', b'N']; 256];
    let mut i = 0;
    while i < 256 {
        let first = match (i >> 4) & 0xF {
            0b1000 => b'A',
            0b0100 => b'C',
            0b0010 => b'G',
            0b0001 => b'T',
            _ => b'N',
        };
        let second = match i & 0xF {
            0b1000 => b'A',
            0b0100 => b'C',
            0b0010 => b'G',
            0b0001 => b'T',
            _ => b'N',
        };
        table[i] = [first, second];
        i += 1;
    }
    table
}

const NUCLEOTIDE_LOOKUP: [[u8; 2]; 256] = create_nucleotide_lookup();

#[pymethods]
impl DnaSequence {
    /// Creates a new DnaSequence from a string of nucleotides
    // #[new]
    // fn new(seq: &str) -> Self {
    //     let num_chunks = (seq.len() + 7) / 8; // Calculate the number of 8-nucleotide chunks
    //     let mut instance = Self {
    //         data: vec![0; num_chunks],
    //         length: seq.len(),
    //     };

    //     let bytes = seq.as_bytes();

    //     for (chunk_idx, chunk) in bytes.chunks(8).enumerate() {
    //         let mut chunk_bits: u32 = 0;
    //         for (i, &nucleotide) in chunk.iter().enumerate() {
    //             let bits = Self::nucleotide_to_onehot_bits(nucleotide.to_ascii_uppercase());
    //             chunk_bits |= (bits as u32) << (28 - 4 * i);
    //         }
    //         instance.data[chunk_idx] = chunk_bits;
    //     }

    //     instance
    // }
    #[new]
    pub fn new(seq: &str) -> Self {
        let num_chunks = (seq.len() + 7) / 8;
        let mut data = vec![0u32; num_chunks];
        let bytes = seq.as_bytes();

        for (chunk_idx, chunk) in bytes.chunks(8).enumerate() {
            let mut chunk_bits: u32 = 0;
            for (i, &nucleotide) in chunk.iter().enumerate() {
                let bits: u32 = match nucleotide {
                    b'A' | b'a' => 0b1000,
                    b'C' | b'c' => 0b0100,
                    b'G' | b'g' => 0b0010,
                    b'T' | b't' => 0b0001,
                    _ => 0b0000,
                };
                chunk_bits |= bits << (28 - 4 * i);
            }
            data[chunk_idx] = chunk_bits;
        }

        Self {
            data,
            length: seq.len(),
        }
    }

    /// Decodes the internal u32 vector back into a string of nucleotides.
    ///
    /// This version uses a lookup table to convert the bits to nucleotides.
    /// This is faster than decoding each 4 bits individually:
    ///
    /// ```no_run,ignore
    ///    fn decode(&self) -> String {
    ///        let mut nucleotides = String::with_capacity(self.length);
    ///        for &chunk in &self.data {
    ///            for i in 0..8 {
    ///                if nucleotides.len() < self.length {
    ///                    let bits = ((chunk >> (28 - 4 * i)) & 0b1111) as u8;
    ///                    nucleotides.push(Self::onehot_bits_to_nucleotide(bits));
    ///                }
    ///            }
    ///        }
    ///     
    ///        nucleotides
    ///   }
    /// ```
    /// 
    pub fn decode(&self) -> String {
        // Pre-allocating full capacity and then truncating is slightly faster
        // than checking the bounds each time
        let capacity = self.data.len() * 8;
        let mut nucleotides = String::with_capacity(capacity);

        for &chunk in &self.data {
            let bytes = chunk.to_be_bytes();
            for &byte in &bytes {
                let [first, second] = NUCLEOTIDE_LOOKUP[byte as usize];
                nucleotides.push(first as char);
                nucleotides.push(second as char);
            }
            if nucleotides.len() == self.length {
                break;
            }
        }

        nucleotides.truncate(self.length);

        nucleotides
    }

    pub fn hamming_distance(&self, other: &Self) -> u32 {
        let min_length = self.length.min(other.length);

        // Iterate over each block
        self.data
            .iter()
            .zip(other.data.iter())
            .enumerate()
            .fold(0, |acc, (idx, (&x, &y))| {
                let block_max_nucleotides = if (idx + 1) * 8 > min_length {
                    min_length % 8
                } else {
                    8
                };

                let valid_bits = if block_max_nucleotides == 0 {
                    8
                } else {
                    block_max_nucleotides
                };

                acc + Self::hamming_distance_u32(x, y, valid_bits)
            })
    }

    /// Returns the data as a NumPy array of one-hot encoded bits.
    pub fn data_as_numpy<'py>(&self, py: Python<'py>) -> &'py PyArray1<u32> {
        // Create a NumPy array from the data vector.
        PyArray::from_slice(py, &self.data)
    }

    /// Returns a view of the data as a NumPy array of u8.
    pub fn data_as_numpy_u8<'py>(&self, py: Python<'py>) -> PyResult<&'py PyArray1<u8>> {
        // Safety: This is safe as long as the data in `self.data` is properly aligned and
        // valid for reinterpretation.
        let data_u8 = unsafe {
            std::slice::from_raw_parts(
                self.data.as_ptr() as *const u8,
                self.data.len() * std::mem::size_of::<u32>(),
            )
        };

        // Create a PyArray1<u8> from the raw bytes.
        Ok(PyArray::from_slice(py, data_u8))
    }

    pub fn count_content(&self, nucleotides: &str) -> u64 {
        if self.length == 0 {
            return 0u64;
        }

        let bases: Vec<char> = nucleotides.chars().collect();

        let mask = Self::create_content_mask(bases.as_slice());
        let mut total_count = 0u64;

        // Handle all complete u32 chunks
        let complete_chunks = self.length / 8;
        for &word in &self.data[..complete_chunks] {
            total_count += (word & mask).count_ones() as u64;
        }

        // Handle the last partial chunk if it exists
        if self.length % 8 != 0 && complete_chunks < self.data.len() {
            let remaining_nucleotides = self.length % 8;
            let last_word = self.data[complete_chunks];
            // Create a mask that only keeps the valid number of 4-bit nucleotides
            let valid_nucleotides_mask = (!0u32) << (32 - (remaining_nucleotides * 4));
            total_count += (last_word & mask & valid_nucleotides_mask).count_ones() as u64;
        }

        total_count
    }

    /// Returns the length of the data
    fn __len__(&self) -> usize {
        self.length
    }

    fn __getitem__(&self, idx: &PyAny) -> PyResult<Self> {
        // Check if idx is a slice
        if let Ok(slice) = idx.downcast::<PySlice>() {
            #[cfg(target_pointer_width = "64")]
            let range = slice.indices(self.length as i64)?;
            
            #[cfg(target_pointer_width = "32")]
            let range = slice.indices((self.length as i64).try_into().unwrap())?;

            // assert step is not different from 1
            if range.step != 1 {
                return Err(PyTypeError::new_err("Slicing with a step is not supported"));
            }

            let start = range.start.clamp(0, self.length as isize) as usize;
            let stop = range.stop.clamp(0, self.length as isize) as usize;

            let length = stop - start;
            let mut data = Vec::with_capacity((length + 7) / 8);

            // Start at the correct byte within the first chunk
            let bit_offset = (start % 8) * 4;

            if bit_offset == 0 {
                let length = (stop - start).clamp(0, self.length);
                let mut data = vec![0u32; (length + 7) / 8];
                data.copy_from_slice(&self.data[start / 8..(stop + 7) / 8]);

                return Ok(Self { data, length });
            }


            // // Repackage the data if the slice starts in the middle of a chunk
            // let mut current_pos = start / 8;
            // while current_pos < (stop + 7) / 8 {
            //     let current_chunk = self.data[current_pos];
            //     let next_chunk = if current_pos + 1 < self.data.len() {
            //         self.data[current_pos + 1]
            //     } else {
            //         0
            //     };

            //     // Combine parts from current and next chunk
            //     let mut new_chunk = current_chunk << bit_offset;
            //     if bit_offset > 0 {
            //         new_chunk |= next_chunk >> (32 - bit_offset);
            //     }
                
            //     data.push(new_chunk);
            //     current_pos += 1;
            // }
            // Repackage the data if the slice starts in the middle of a chunk
            let mut chunk: u32;
            let mut next_chunk: u32;
            for i in start / 8..(stop + 7) / 8 {
                chunk = self.data[i];
                if i + 1 < self.data.len() {
                    next_chunk = self.data[i + 1];
                } else {
                    next_chunk = 0;
                }

                // Extract the desired bits, shifting right if necessary
                let nucleotides = chunk << bit_offset;
                let nucleotides = nucleotides | (next_chunk >> (32 - bit_offset));
                data.push(nucleotides);
            }

            Ok(Self { data, length })
        } else if let Ok(index) = idx.extract::<usize>() {
            let bit_offset = (index % 8) * 4;
            if index < self.length {
                let mut data = Vec::with_capacity(1);
                data.push((self.data[index / 8] >> (28 - bit_offset) & 0b1111) << 28);
                Ok(Self { data, length: 1 })
            } else {
                Err(PyTypeError::new_err("Index out of bounds"))
            }
        } else {
            Err(PyTypeError::new_err("Index must be a slice"))
        }
    }

    /// Involved compute - decoding the data
    fn __str__(&self) -> String {
        self.decode()
    }

    /// Unpacks the DNA sequence data into a bit representation
    /// 
    /// NumPy version:
    /// ```no_run,ignore
    ///     np.unpackbits(
    ///         obj.data_as_numpy().byteswap().view(np.uint8).reshape(-1, 4), 
    ///         count=len(obj) * 4
    ///     ).reshape(-1, 4)
    /// ```
    pub fn unpack_bits<'py>(&self, py: Python<'py>) -> PyResult<&'py PyArray2<u8>> {
        // Get the data as numpy array
        let numpy_data = self.data_as_numpy(py);
        
        // Get numpy module and uint8 dtype
        let np = py.import("numpy")?;
        let uint8 = np.getattr("uint8")?;
        
        // Byteswap the data
        let byteswapped = numpy_data.call_method0("byteswap")?;
        
        // Convert to u8 view with explicit dtype
        let view_u8 = byteswapped.call_method1("view", (uint8,))?;
        
        // Use numpy's unpackbits function
        let kwargs = pyo3::types::PyDict::new(py);
        kwargs.set_item("count", self.length * 4)?;
        
        let bits = np.getattr("unpackbits")?.call((view_u8,), Some(kwargs))?;
        
        // Reshape to (sequence_length, 4)
        let reshaped = bits.call_method1("reshape", (self.length, 4))?;
        
        // Convert back to Rust PyArray
        Ok(reshaped.extract()?)
    }

    fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<DnaSequenceIterator>> {
        let iter = DnaSequenceIterator::new(slf.data.clone(), slf.length);
        Py::new(slf.py(), iter)
    }
}

impl DnaSequence {
    // /// Converts a nucleotide character to its corresponding bits.
    // fn nucleotide_to_onehot_bits(n: u8) -> u8 {
    //     match n {
    //         b'A' => 0b1000,
    //         b'C' => 0b0100,
    //         b'G' => 0b0010,
    //         b'T' => 0b0001,
    //         _ => 0b0000,
    //     }
    // }

    // /// Converts two bits back into a nucleotide character.
    // fn onehot_bits_to_nucleotide(bits: u8) -> char {
    //     match bits {
    //         0b1000 => 'A',
    //         0b0100 => 'C',
    //         0b0010 => 'G',
    //         0b0001 => 'T',
    //         _ => 'N',
    //         // _ => panic!("Invalid bits: {:b}", bits),
    //     }
    // }

    fn hamming_distance_u32(a: u32, b: u32, valid_bits: usize) -> u32 {
        let diff = a ^ b;
        let mask = if valid_bits == 8 {
            0xFFFFFFFF
        } else {
            // create a mask and align it to 32 bits
            ((1 << (4 * valid_bits)) - 1) << (32 - 4 * valid_bits)
        };

        let masked_diff = diff & mask;

        // Apply mask trick to count the number of differing nucleotides
        let mask1 = 0x11111111; // Mask to isolate the least significant bit of each 4-bit nucleotide
        let mask2 = 0x22222222; // Mask to isolate the second least significant bit of each 4-bit nucleotide
        let mask3 = 0x44444444; // Mask to isolate the third least significant bit of each 4-bit nucleotide
        let mask4 = 0x88888888; // Mask to isolate the most significant bit of each 4-bit nucleotide

        let diff1 = masked_diff & mask1;
        let diff2 = (masked_diff & mask2) >> 1;
        let diff3 = (masked_diff & mask3) >> 2;
        let diff4 = (masked_diff & mask4) >> 3;

        let out = diff1 | diff2 | diff3 | diff4;

        out.count_ones()
    }

    const AAAA: u32 = 0b1000_1000_1000_1000_1000_1000_1000_1000;
    const CCCC: u32 = 0b0100_0100_0100_0100_0100_0100_0100_0100;
    const GGGG: u32 = 0b0010_0010_0010_0010_0010_0010_0010_0010;
    const TTTT: u32 = 0b0001_0001_0001_0001_0001_0001_0001_0001;
    const NNNN: u32 = 0b0000_0000_0000_0000_0000_0000_0000_0000;

    fn create_content_mask(nucleotides: &[char]) -> u32 {
        let mut mask = 0u32;
        for &c in nucleotides {
            let pattern = match c {
                'A' => Self::AAAA,
                'C' => Self::CCCC,
                'G' => Self::GGGG,
                'T' => Self::TTTT,
                _ => Self::NNNN,
            };
            mask |= pattern;
        }
        mask
    }
}

#[pyclass]
pub struct DnaSequenceIterator {
    data: Vec<u32>,
    length: usize,
    position: usize,
}

#[pymethods]
impl DnaSequenceIterator {
    #[new]
    fn new(data: Vec<u32>, length: usize) -> Self {
        Self {
            data,
            length,
            position: 0,
        }
    }

    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __next__(&mut self) -> Option<char> {
        if self.position >= self.length {
            return None;
        }

        let chunk_idx = self.position / 8;
        let pos_in_chunk = self.position % 8;
        let chunk = self.data[chunk_idx];
        
        let bits = (chunk >> (28 - 4 * pos_in_chunk)) & 0b1111;
        
        let nucleotide = match bits {
            0b1000 => 'A',
            0b0100 => 'C',
            0b0010 => 'G',
            0b0001 => 'T',
            _ => 'N',
        };

        self.position += 1;
        Some(nucleotide)
    }
}

// For Rust-side iteration, we can create a separate iterator type
pub struct RustDnaIterator<'a> {
    sequence: &'a DnaSequence,
    position: usize,
}

impl<'a> Iterator for RustDnaIterator<'a> {
    type Item = char;

    fn next(&mut self) -> Option<Self::Item> {
        if self.position >= self.sequence.length {
            return None;
        }

        let chunk_idx = self.position / 8;
        let pos_in_chunk = self.position % 8;
        let chunk = self.sequence.data[chunk_idx];
        
        let bits = (chunk >> (28 - 4 * pos_in_chunk)) & 0b1111;
        
        let nucleotide = match bits {
            0b1000 => 'A',
            0b0100 => 'C',
            0b0010 => 'G',
            0b0001 => 'T',
            _ => 'N',
        };

        self.position += 1;
        Some(nucleotide)
    }
}

impl DnaSequence {
    // Add a method for Rust-side iteration
    pub fn iter(&self) -> RustDnaIterator {
        RustDnaIterator {
            sequence: self,
            position: 0,
        }
    }
}

impl<'a> IntoIterator for &'a DnaSequence {
    type Item = char;
    type IntoIter = RustDnaIterator<'a>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

/// A Python module implemented in Rust.
#[pymodule]
fn ganz(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    m.add_class::<DnaSequence>()?;
    Ok(())
}
