from django import template
from iching.utils import bz

register = template.Library()

@register.filter
def get_month_earth_branch(solar_term_index):
    """
    Convert solar term index to month's earth branch.
    Solar terms come in pairs for each month:
    0,1 = 寅 (first month)
    2,3 = 卯 (second month)
    ...and so on
    """
    if solar_term_index is None:
        return ''
        
    try:
        # Calculate month index (0-11) from solar term index
        month_index = int(solar_term_index) // 2
        
        # Map to earth branches starting from 寅
        earth_branches = ['寅', '卯', '辰', '巳', '午', '未', '申', '酉', '戌', '亥', '子', '丑']
        return earth_branches[month_index % 12]
    except (ValueError, TypeError):
        return ''

@register.filter
def earth_to_element(earth_branch):
    """Convert earth branch to its corresponding element."""
    # Earth branch to index mapping
    earth_to_index = {
        '子': 0, '丑': 1, '寅': 2, '卯': 3, '辰': 4, '巳': 5,
        '午': 6, '未': 7, '申': 8, '酉': 9, '戌': 10, '亥': 11
    }
    
    # Get the index for the earth branch
    index = earth_to_index.get(earth_branch)
    if index is None:
        return ''
    
    # Map earth branch index to element
    # 寅卯 (2,3) -> wood
    # 巳午 (5,6) -> fire
    # 辰未戌丑 (4,7,10,1) -> earth
    # 申酉 (8,9) -> metal
    # 亥子 (11,0) -> water
    if index in [2, 3]:
        return 'wood'
    elif index in [5, 6]:
        return 'fire'
    elif index in [4, 7, 10, 1]:
        return 'earth'
    elif index in [8, 9]:
        return 'metal'
    elif index in [11, 0]:
        return 'water'
    return ''

@register.filter
def get_month_element(solar_term_index):
    """Get the element for a month based on its solar term index."""
    earth_branch = get_month_earth_branch(solar_term_index)
    return earth_to_element(earth_branch) 